From 018c75c918a2df4a698a09e882bf4689bef6f00f Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Tue, 16 Jun 2026 23:19:06 -0700 Subject: [PATCH 1/8] exec-server: add remote environment connection lifecycle --- codex-rs/exec-server/src/client.rs | 343 ++++++++++++++++++++++-- codex-rs/exec-server/src/environment.rs | 75 +++--- codex-rs/exec-server/tests/relay.rs | 52 ++-- 3 files changed, 379 insertions(+), 91 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index ae77dc7c4fe..5dcaafa130c 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -14,7 +14,7 @@ use futures::FutureExt; use futures::future::BoxFuture; use serde_json::Value; use tokio::sync::Mutex; -use tokio::sync::Semaphore; +use tokio::sync::OnceCell; use tokio::sync::mpsc; use tokio::sync::watch; @@ -107,6 +107,13 @@ const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); const PROCESS_EVENT_CHANNEL_CAPACITY: usize = 256; const PROCESS_EVENT_RETAINED_BYTES: usize = 1024 * 1024; +const ENVIRONMENT_STARTUP_TIMEOUT: Duration = Duration::from_secs(5 * 60); +const ENVIRONMENT_INITIAL_RETRY_DELAY: Duration = Duration::from_secs(1); +const ENVIRONMENT_MAX_RETRY_DELAY: Duration = Duration::from_secs(30); + +// ThreadEnvironments::snapshot() currently waits for Environment::info(), so this +// must land with the follow-up that stops snapshots from waiting on starting environments. +// Otherwise an unavailable executor can delay session startup for five minutes. impl Default for ExecServerClientConnectOptions { fn default() -> Self { @@ -226,54 +233,125 @@ impl Drop for ActiveProcessStart { } } +type ConnectionResult = Result>; +type ConnectionAttempt = OnceCell; + #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { transport_params: ExecServerTransportParams, - client: Arc>>, - connect_lock: Arc, + // Saves the first startup result so callers share it and failures remain final. + startup: Arc, + // The latest successful client, replaced whenever reconnecting succeeds. + current_client: Arc>>, + reconnect: Arc>>>, } impl LazyRemoteExecServerClient { pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self { Self { transport_params, - client: Arc::new(StdMutex::new(None)), - connect_lock: Arc::new(Semaphore::new(/*permits*/ 1)), + startup: Arc::new(ConnectionAttempt::new()), + current_client: Arc::new(StdMutex::new(None)), + reconnect: Arc::new(StdMutex::new(None)), } } + pub(crate) fn start_connecting(&self) { + let client = self.clone(); + drop(tokio::spawn(async move { + if let Err(error) = client.wait_until_ready().await { + debug!(%error, "exec-server environment startup failed"); + } + })); + } + + pub(crate) fn startup_finished(&self) -> bool { + self.startup.get().is_some() + } + + pub(crate) async fn wait_until_ready(&self) -> Result<(), ExecServerError> { + self.initial_client().await.map(drop) + } + pub(crate) async fn get(&self) -> Result { if let Some(client) = self.connected_client() { return Ok(client); } - let _connect_permit = self.connect_lock.acquire().await.map_err(|_| { - ExecServerError::Protocol("exec-server connect lock closed".to_string()) - })?; - if let Some(client) = self.connected_client() { - return Ok(client); + let Some(cached_client) = self.cached_client() else { + let client = self.initial_client().await?; + if !client.is_disconnected() || !self.can_reconnect() { + return Ok(client); + } + return self.reconnect().await; + }; + + if !self.can_reconnect() { + return Ok(cached_client); } - let next_client = match self.cached_client() { - Some(_client) - if matches!( - &self.transport_params, - ExecServerTransportParams::WebSocketUrl { .. } - | ExecServerTransportParams::NoiseRendezvous { .. } - ) => - { - ExecServerClient::connect_for_transport(self.transport_params.clone()).await? + self.reconnect().await + } + + async fn initial_client(&self) -> Result { + // The first caller starts the work; every other caller waits for that same result. + let result = self + .startup + .get_or_init(|| connect_with_startup_retries(self.transport_params.clone())) + .await; + match result { + Ok(client) => { + let mut current_client = self + .current_client + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if current_client.is_none() { + *current_client = Some(client.clone()); + } + Ok(client.clone()) } - Some(client) => return Ok(client), - None => ExecServerClient::connect_for_transport(self.transport_params.clone()).await?, - }; + Err(error) => Err(ExecServerError::ConnectionAttempt(Arc::clone(error))), + } + } - let mut cached_client = self - .client + async fn reconnect(&self) -> Result { + // Callers handling the same outage share one reconnect attempt. + let attempt = { + let mut reconnect = self + .reconnect + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(client) = self.connected_client() { + return Ok(client); + } + reconnect + .get_or_insert_with(|| Arc::new(ConnectionAttempt::new())) + .clone() + }; + let result = attempt + .get_or_init(|| async { + let result = connect_once(self.transport_params.clone()).await; + if let Ok(client) = &result { + *self + .current_client + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(client.clone()); + } + result + }) + .await; + let mut reconnect = self + .reconnect .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - *cached_client = Some(next_client.clone()); - Ok(next_client) + // Forget only this completed attempt so a later operation can retry after failure. + if reconnect + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &attempt)) + { + *reconnect = None; + } + result.clone().map_err(ExecServerError::ConnectionAttempt) } fn connected_client(&self) -> Option { @@ -282,11 +360,60 @@ impl LazyRemoteExecServerClient { } fn cached_client(&self) -> Option { - self.client + self.current_client .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .clone() } + + fn can_reconnect(&self) -> bool { + matches!( + self.transport_params, + ExecServerTransportParams::WebSocketUrl { .. } + | ExecServerTransportParams::NoiseRendezvous { .. } + ) + } +} + +async fn connect_with_startup_retries( + transport_params: ExecServerTransportParams, +) -> ConnectionResult { + if matches!( + transport_params, + ExecServerTransportParams::StdioCommand { .. } + ) { + return connect_once(transport_params).await; + } + + let startup = async { + let mut retry_delay = ENVIRONMENT_INITIAL_RETRY_DELAY; + loop { + match ExecServerClient::connect_for_transport(transport_params.clone()).await { + Ok(client) => return Ok(client), + Err(error) => { + debug!( + %error, + retry_in = ?retry_delay, + "exec-server environment is not ready; retrying" + ); + tokio::time::sleep(retry_delay).await; + retry_delay = (retry_delay * 2).min(ENVIRONMENT_MAX_RETRY_DELAY); + } + } + } + }; + match timeout(ENVIRONMENT_STARTUP_TIMEOUT, startup).await { + Ok(result) => result, + Err(_) => Err(Arc::new(ExecServerError::StartupTimedOut { + timeout: ENVIRONMENT_STARTUP_TIMEOUT, + })), + } +} + +async fn connect_once(transport_params: ExecServerTransportParams) -> ConnectionResult { + ExecServerClient::connect_for_transport(transport_params) + .await + .map_err(Arc::new) } impl HttpClient for LazyRemoteExecServerClient { @@ -352,6 +479,10 @@ pub enum ExecServerError { EnvironmentRegistryAuth(String), #[error("environment registry request failed: {0}")] EnvironmentRegistryRequest(#[from] reqwest::Error), + #[error("exec-server connection attempt failed: {0}")] + ConnectionAttempt(#[source] Arc), + #[error("exec-server did not become ready within {timeout:?}")] + StartupTimedOut { timeout: Duration }, } impl ExecServerClient { @@ -1773,6 +1904,164 @@ mod tests { server.await.expect("server task should finish"); } + #[tokio::test] + async fn initial_connection_retries_once_and_is_shared_by_all_waiters() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let server = tokio::spawn(async move { + let (first, _) = listener + .accept() + .await + .expect("first connection should arrive"); + drop(first); + + let mut second = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut second, + "startup-session", + /*expected_resume_session_id*/ None, + ) + .await; + timeout(Duration::from_secs(1), second.next()) + .await + .expect("client should close after the test"); + }); + let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }); + + assert!(!client.startup_finished()); + client.start_connecting(); + let (ready, first, second) = + tokio::join!(client.wait_until_ready(), client.get(), client.get()); + ready.expect("background startup should finish"); + let first = first.expect("first waiter should receive the client"); + let second = second.expect("second waiter should receive the same client"); + + assert!(client.startup_finished()); + assert_eq!(first.session_id().as_deref(), Some("startup-session")); + assert!(Arc::ptr_eq(&first.inner, &second.inner)); + + drop(first); + drop(second); + drop(client); + server.await.expect("server task should finish"); + } + + #[tokio::test] + async fn terminal_stdio_startup_failure_is_remembered() { + let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::StdioCommand { + command: StdioExecServerCommand { + program: "codex-missing-exec-server-for-test".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: None, + }, + initialize_timeout: Duration::from_secs(1), + }); + + let first = match client.get().await { + Ok(_) => panic!("missing executable should fail"), + Err(error) => error, + }; + assert!(client.startup_finished()); + let second = match client.get().await { + Ok(_) => panic!("burned environment should stay failed"), + Err(error) => error, + }; + + let ( + super::ExecServerError::ConnectionAttempt(first), + super::ExecServerError::ConnectionAttempt(second), + ) = (first, second) + else { + panic!("expected saved connection failures"); + }; + assert!(Arc::ptr_eq(&first, &second)); + } + + #[tokio::test] + async fn failed_reconnect_does_not_burn_environment() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let (replacement_initialized_tx, replacement_initialized_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let mut first = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut first, + "startup-session", + /*expected_resume_session_id*/ None, + ) + .await; + first + .close(None) + .await + .expect("startup websocket should close"); + + let (mut failed_reconnect, _) = listener + .accept() + .await + .expect("first reconnect should arrive"); + failed_reconnect + .write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n") + .await + .expect("failed handshake response should write"); + drop(failed_reconnect); + + let mut successful_reconnect = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut successful_reconnect, + "replacement-session", + /*expected_resume_session_id*/ None, + ) + .await; + replacement_initialized_tx + .send(()) + .expect("replacement initialization should be observed"); + timeout(Duration::from_secs(1), successful_reconnect.next()) + .await + .expect("client should close after the test"); + }); + let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }); + + let initial = client.get().await.expect("startup should connect"); + wait_for_disconnect(&initial).await; + assert!(matches!( + client.get().await, + Err(super::ExecServerError::ConnectionAttempt(_)) + )); + let replacement = client.get().await.expect("later reconnect should succeed"); + + assert_eq!( + replacement.session_id().as_deref(), + Some("replacement-session") + ); + replacement_initialized_rx + .await + .expect("server should observe replacement initialization"); + + drop(initial); + drop(replacement); + drop(client); + server.await.expect("server task should finish"); + } + #[tokio::test] async fn wake_notifications_do_not_block_other_sessions() { let (client_stdin, server_reader) = duplex(1 << 20); diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 6159fb2dd26..76f5d7b2c9b 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -2,9 +2,6 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::RwLock; -use futures::FutureExt; -use futures::future::BoxFuture; - use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::ExecutorFileSystem; @@ -402,50 +399,19 @@ fn optional_environment_value(name: &str) -> Option { #[derive(Clone)] pub struct Environment { exec_server_url: Option, - remote_transport: Option, - info_provider: Arc, + remote_client: Option, exec_backend: Arc, filesystem: Arc, http_client: Arc, local_runtime_paths: Option, } -/// Provides environment metadata from either a local environment or a remote exec-server. -trait EnvironmentInfoProvider: Send + Sync { - fn info(&self) -> BoxFuture<'_, Result>; -} - -struct LocalEnvironmentInfoProvider; - -impl EnvironmentInfoProvider for LocalEnvironmentInfoProvider { - fn info(&self) -> BoxFuture<'_, Result> { - std::future::ready(Ok(EnvironmentInfo::local())).boxed() - } -} - -struct RemoteEnvironmentInfoProvider { - client: LazyRemoteExecServerClient, -} - -impl RemoteEnvironmentInfoProvider { - fn new(client: LazyRemoteExecServerClient) -> Self { - Self { client } - } -} - -impl EnvironmentInfoProvider for RemoteEnvironmentInfoProvider { - fn info(&self) -> BoxFuture<'_, Result> { - async move { self.client.environment_info().await }.boxed() - } -} - impl Environment { /// Builds a test-only local environment without configured sandbox helper paths. pub fn default_for_tests() -> Self { Self { exec_server_url: None, - remote_transport: None, - info_provider: Arc::new(LocalEnvironmentInfoProvider), + remote_client: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), http_client: Arc::new(ReqwestHttpClient), @@ -501,8 +467,7 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { exec_server_url: None, - remote_transport: None, - info_provider: Arc::new(LocalEnvironmentInfoProvider), + remote_client: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::with_runtime_paths( local_runtime_paths.clone(), @@ -534,15 +499,14 @@ impl Environment { ExecServerTransportParams::NoiseRendezvous { .. } => None, ExecServerTransportParams::StdioCommand { .. } => None, }; - let client = LazyRemoteExecServerClient::new(remote_transport.clone()); + let client = LazyRemoteExecServerClient::new(remote_transport); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); Self { exec_server_url, - remote_transport: Some(remote_transport), - info_provider: Arc::new(RemoteEnvironmentInfoProvider::new(client.clone())), + remote_client: Some(client.clone()), exec_backend, filesystem, http_client: Arc::new(client), @@ -551,7 +515,7 @@ impl Environment { } pub fn is_remote(&self) -> bool { - self.remote_transport.is_some() + self.remote_client.is_some() } /// Returns the remote exec-server URL when this environment is remote. @@ -565,7 +529,32 @@ impl Environment { /// Returns environment information from the selected execution/filesystem environment. pub async fn info(&self) -> Result { - self.info_provider.info().await + match &self.remote_client { + Some(client) => client.environment_info().await, + None => Ok(EnvironmentInfo::local()), + } + } + + /// Starts connecting a remote environment without waiting for it. + pub fn start_connecting(&self) { + if let Some(client) = &self.remote_client { + client.start_connecting(); + } + } + + /// Returns whether initial startup has either succeeded or permanently failed. + pub fn startup_finished(&self) -> bool { + self.remote_client + .as_ref() + .is_none_or(LazyRemoteExecServerClient::startup_finished) + } + + /// Waits for initial startup. A failed startup is never attempted again. + pub async fn wait_until_ready(&self) -> Result<(), ExecServerError> { + match &self.remote_client { + Some(client) => client.wait_until_ready().await, + None => Ok(()), + } } pub fn get_exec_backend(&self) -> Arc { diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 114fe122a25..cbb0259ec75 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -6,8 +6,6 @@ mod relay_proto; use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Context; @@ -43,6 +41,7 @@ use relay_proto::relay_message_frame; use tempfile::TempDir; use tokio::net::TcpListener; use tokio::net::TcpStream; +use tokio::sync::mpsc; use tokio::time::timeout; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::accept_async; @@ -73,7 +72,7 @@ impl AuthProvider for StaticRegistryAuthProvider { } struct FailingNoiseConnectProvider { - attempts: Arc, + attempt_tx: mpsc::UnboundedSender<()>, } impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { @@ -81,7 +80,7 @@ impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { &self, _: NoiseChannelPublicKey, ) -> BoxFuture<'_, Result> { - self.attempts.fetch_add(1, Ordering::SeqCst); + let _ = self.attempt_tx.send(()); async { Err(ExecServerError::Protocol( "test registry connect failure".to_string(), @@ -96,41 +95,52 @@ fn static_registry_auth_provider() -> codex_api::SharedAuthProvider { } #[tokio::test] -async fn noise_environment_refreshes_bundle_for_each_connection_attempt() -> Result<()> { - let attempts = Arc::new(AtomicUsize::new(0)); +async fn noise_environment_refreshes_bundle_during_startup_retries() -> Result<()> { + let (attempt_tx, mut attempt_rx) = mpsc::unbounded_channel(); let manager = EnvironmentManager::without_environments(); manager.upsert_noise_environment( ENVIRONMENT_ID.to_string(), - Arc::new(FailingNoiseConnectProvider { - attempts: Arc::clone(&attempts), - }), + Arc::new(FailingNoiseConnectProvider { attempt_tx }), )?; let backend = manager .get_environment(ENVIRONMENT_ID) .context("Noise environment should be materialized")? .get_exec_backend(); + let cwd = PathUri::from_path(std::env::current_dir()?)?; - for attempt in 1..=2 { - let result = backend + let startup = tokio::spawn(async move { + backend .start(ExecParams { - process_id: ProcessId::new(format!("proc-{attempt}")), + process_id: ProcessId::from("proc-1"), argv: vec!["true".to_string()], - cwd: PathUri::from_path(std::env::current_dir()?)?, + cwd, env_policy: None, env: HashMap::new(), tty: false, pipe_stdin: false, arg0: None, }) - .await; - assert!(matches!( - result, - Err(ExecServerError::Protocol(ref message)) - if message == "test registry connect failure" - )); - } + .await + }); + + timeout(TEST_TIMEOUT, async { + for _ in 0..2 { + attempt_rx + .recv() + .await + .context("connection provider should remain available")?; + } + Ok::<_, anyhow::Error>(()) + }) + .await + .context("startup should retry the Noise connection")??; - assert_eq!(attempts.load(Ordering::SeqCst), 2); + startup.abort(); + let cancellation = match startup.await { + Err(error) => error, + Ok(_) => panic!("startup should be cancelled"), + }; + assert!(cancellation.is_cancelled()); Ok(()) } From bd9f20232f1c06b5a91cf2c3f8330d1e578fcc3a Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Tue, 16 Jun 2026 23:19:57 -0700 Subject: [PATCH 2/8] exec-server: remove temporary landing note --- codex-rs/exec-server/src/client.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 5dcaafa130c..ca6b103c89b 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -111,10 +111,6 @@ const ENVIRONMENT_STARTUP_TIMEOUT: Duration = Duration::from_secs(5 * 60); const ENVIRONMENT_INITIAL_RETRY_DELAY: Duration = Duration::from_secs(1); const ENVIRONMENT_MAX_RETRY_DELAY: Duration = Duration::from_secs(30); -// ThreadEnvironments::snapshot() currently waits for Environment::info(), so this -// must land with the follow-up that stops snapshots from waiting on starting environments. -// Otherwise an unavailable executor can delay session startup for five minutes. - impl Default for ExecServerClientConnectOptions { fn default() -> Self { Self { From c171ed3ac0503681e861fc7540cb03947be89489 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Tue, 16 Jun 2026 23:37:00 -0700 Subject: [PATCH 3/8] exec-server: fail fast on permanent startup errors --- codex-rs/exec-server/src/client.rs | 56 ++++++++++++++++++++++++++++- codex-rs/exec-server/tests/relay.rs | 9 +++-- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index ca6b103c89b..30e0185a7ab 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -386,7 +386,7 @@ async fn connect_with_startup_retries( loop { match ExecServerClient::connect_for_transport(transport_params.clone()).await { Ok(client) => return Ok(client), - Err(error) => { + Err(error) if error.is_retryable_during_startup() => { debug!( %error, retry_in = ?retry_delay, @@ -395,6 +395,7 @@ async fn connect_with_startup_retries( tokio::time::sleep(retry_delay).await; retry_delay = (retry_delay * 2).min(ENVIRONMENT_MAX_RETRY_DELAY); } + Err(error) => return Err(Arc::new(error)), } } }; @@ -481,6 +482,36 @@ pub enum ExecServerError { StartupTimedOut { timeout: Duration }, } +impl ExecServerError { + fn is_retryable_during_startup(&self) -> bool { + // This is a positive allowlist so new error variants fail fast by default. + match self { + Self::WebSocketConnectTimeout { .. } + | Self::InitializeTimedOut { .. } + | Self::Closed + | Self::Disconnected(_) => true, + Self::WebSocketConnect { source, .. } => match source { + tokio_tungstenite::tungstenite::Error::Io(_) => true, + tokio_tungstenite::tungstenite::Error::Http(response) => { + retryable_http_status(response.status()) + } + _ => false, + }, + Self::EnvironmentRegistryHttp { status, .. } => retryable_http_status(*status), + Self::EnvironmentRegistryRequest(error) => error.is_connect() || error.is_timeout(), + _ => false, + } + } +} + +fn retryable_http_status(status: reqwest::StatusCode) -> bool { + status.is_server_error() + || matches!( + status, + reqwest::StatusCode::REQUEST_TIMEOUT | reqwest::StatusCode::TOO_MANY_REQUESTS + ) +} + impl ExecServerClient { pub async fn initialize( &self, @@ -1900,6 +1931,29 @@ mod tests { server.await.expect("server task should finish"); } + #[tokio::test] + async fn invalid_websocket_url_fails_without_startup_retries() { + let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { + websocket_url: "http://example.com".to_string(), + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }); + + let result = timeout(Duration::from_secs(1), client.get()) + .await + .expect("permanent startup failure should return immediately"); + let error = match result { + Ok(_) => panic!("invalid URL should fail"), + Err(error) => error, + }; + + assert!(matches!( + error, + super::ExecServerError::ConnectionAttempt(source) + if matches!(source.as_ref(), super::ExecServerError::WebSocketConnect { .. }) + )); + } + #[tokio::test] async fn initial_connection_retries_once_and_is_shared_by_all_waiters() { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index cbb0259ec75..6d3cc3e53ea 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -34,6 +34,7 @@ use futures::StreamExt; use futures::future::BoxFuture; use http::HeaderMap; use http::HeaderValue; +use http::StatusCode; use pretty_assertions::assert_eq; use prost::Message as ProstMessage; use relay_proto::RelayMessageFrame; @@ -82,9 +83,11 @@ impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { ) -> BoxFuture<'_, Result> { let _ = self.attempt_tx.send(()); async { - Err(ExecServerError::Protocol( - "test registry connect failure".to_string(), - )) + Err(ExecServerError::EnvironmentRegistryHttp { + status: StatusCode::SERVICE_UNAVAILABLE, + code: None, + message: "test registry unavailable".to_string(), + }) } .boxed() } From f4f57bb882a8886b2b7e84892839ce8ed74cca60 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Wed, 17 Jun 2026 18:06:42 -0700 Subject: [PATCH 4/8] exec-server: start remote environments on registration --- codex-rs/exec-server/src/client.rs | 40 ++++++++++++------- codex-rs/exec-server/src/environment.rs | 53 ++++++++++++++++++++----- 2 files changed, 67 insertions(+), 26 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 30e0185a7ab..f2f6bb050cd 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -1245,6 +1245,7 @@ mod tests { use tokio::net::TcpStream; use tokio::sync::mpsc; use tokio::sync::oneshot; + use tokio::sync::watch; use tokio::time::Duration; #[cfg(unix)] use tokio::time::sleep; @@ -2047,6 +2048,7 @@ mod tests { listener.local_addr().expect("listener should have address") ); let (replacement_initialized_tx, replacement_initialized_rx) = oneshot::channel(); + let (allow_replacement_tx, allow_replacement_rx) = watch::channel(false); let server = tokio::spawn(async move { let mut first = accept_websocket(&listener).await; complete_websocket_initialize( @@ -2060,17 +2062,20 @@ mod tests { .await .expect("startup websocket should close"); - let (mut failed_reconnect, _) = listener - .accept() - .await - .expect("first reconnect should arrive"); - failed_reconnect - .write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n") + let successful_reconnect = loop { + let (stream, _) = listener.accept().await.expect("reconnect should arrive"); + if *allow_replacement_rx.borrow() { + break stream; + } + let mut failed_reconnect = stream; + failed_reconnect + .write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n") + .await + .expect("failed handshake response should write"); + }; + let mut successful_reconnect = accept_async(successful_reconnect) .await - .expect("failed handshake response should write"); - drop(failed_reconnect); - - let mut successful_reconnect = accept_websocket(&listener).await; + .expect("replacement websocket handshake should succeed"); complete_websocket_initialize( &mut successful_reconnect, "replacement-session", @@ -2091,11 +2096,16 @@ mod tests { }); let initial = client.get().await.expect("startup should connect"); - wait_for_disconnect(&initial).await; - assert!(matches!( - client.get().await, - Err(super::ExecServerError::ConnectionAttempt(_)) - )); + timeout(Duration::from_secs(1), async { + while !initial.is_disconnected() { + tokio::task::yield_now().await; + } + }) + .await + .expect("client should observe disconnect"); + allow_replacement_tx + .send(true) + .expect("server should allow a fresh client"); let replacement = client.get().await.expect("later reconnect should succeed"); assert_eq!( diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 76f5d7b2c9b..1d77b0c7ae5 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -47,9 +47,9 @@ pub const CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR: &str = /// use `default_environment().is_some()` as the signal for model-facing /// shell/filesystem tool availability. /// -/// Remote environments create remote filesystem and execution backends that -/// lazy-connect to the configured exec-server on first use. The remote -/// transport is not opened when the manager or environment is constructed. +/// Remote environments begin connecting when added to the manager. Their +/// filesystem and execution backends share that startup result and reconnect +/// after later disconnects as needed. #[derive(Debug)] pub struct EnvironmentManager { default_environment: Option, @@ -221,6 +221,10 @@ impl EnvironmentManager { Some(environment_id) } }; + // The snapshot is valid; start connecting its remote environments in the background. + for environment in environment_map.values() { + environment.start_connecting(); + } Ok(Self { default_environment, environments: RwLock::new(environment_map), @@ -304,12 +308,15 @@ impl EnvironmentManager { "remote environment requires an exec-server url".to_string(), )); }; - let environment = - Environment::remote_inner(exec_server_url, self.local_runtime_paths.clone()); + let environment = Arc::new(Environment::remote_inner( + exec_server_url, + self.local_runtime_paths.clone(), + )); + environment.start_connecting(); self.environments .write() .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(environment_id, Arc::new(environment)); + .insert(environment_id, environment); Ok(()) } @@ -333,14 +340,15 @@ impl EnvironmentManager { "failed to generate Noise harness identity: {error}" )) })?; - let environment = Environment::remote_with_transport( + let environment = Arc::new(Environment::remote_with_transport( ExecServerTransportParams::NoiseRendezvous { provider, identity }, self.local_runtime_paths.clone(), - ); + )); + environment.start_connecting(); self.environments .write() .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(environment_id, Arc::new(environment)); + .insert(environment_id, environment); Ok(()) } } @@ -590,6 +598,7 @@ impl From for ShellInfo { #[cfg(test)] mod tests { use std::sync::Arc; + use std::time::Duration; use super::Environment; use super::EnvironmentManager; @@ -602,6 +611,8 @@ mod tests { use crate::environment_provider::EnvironmentProviderSnapshot; use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; + use tokio::net::TcpListener; + use tokio::time::timeout; fn test_runtime_paths() -> ExecServerRuntimePaths { ExecServerRuntimePaths::new( @@ -615,8 +626,8 @@ mod tests { assert!(manager.try_local_environment().is_none()); } - #[test] - fn noise_environment_config_selects_remote_as_default() { + #[tokio::test] + async fn noise_environment_config_selects_remote_as_default() { let config = noise_environment_config_from_values( Some("http://registry.example/api".to_string()), Some("environment-requested".to_string()), @@ -971,6 +982,26 @@ mod tests { assert!(!Arc::ptr_eq(&first, &second)); } + #[tokio::test] + async fn environment_manager_starts_remote_environment_when_upserted() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind websocket listener"); + let manager = EnvironmentManager::without_environments(); + + manager + .upsert_environment( + "executor-a".to_string(), + format!("ws://{}", listener.local_addr().expect("listener address")), + ) + .expect("remote environment"); + + timeout(Duration::from_secs(5), listener.accept()) + .await + .expect("environment should start connecting when registered") + .expect("accept connection"); + } + #[tokio::test] async fn environment_manager_rejects_empty_remote_environment_url() { let manager = EnvironmentManager::without_environments(); From e4f8441cab402d71fbe6e5ef800ef7f05b2f4012 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Wed, 17 Jun 2026 00:27:23 -0700 Subject: [PATCH 5/8] core: track starting environments in snapshots --- codex-rs/core/src/agents_md_tests.rs | 6 +- codex-rs/core/src/environment_selection.rs | 482 +++++++++++++++------ 2 files changed, 358 insertions(+), 130 deletions(-) diff --git a/codex-rs/core/src/agents_md_tests.rs b/codex-rs/core/src/agents_md_tests.rs index 42604629a04..ce0984393b6 100644 --- a/codex-rs/core/src/agents_md_tests.rs +++ b/codex-rs/core/src/agents_md_tests.rs @@ -257,8 +257,8 @@ async fn agents_md_paths(config: &TestConfig) -> std::io::Result( environments: [(&str, AbsolutePathBuf); N], ) -> TurnEnvironmentSnapshot { - TurnEnvironmentSnapshot { - turn_environments: environments + TurnEnvironmentSnapshot::from_turn_environments( + environments .into_iter() .map(|(environment_id, cwd)| { TurnEnvironment::new( @@ -272,7 +272,7 @@ fn resolved_local_environments( ) }) .collect(), - } + ) } fn project_provenance(path: AbsolutePathBuf, cwd: AbsolutePathBuf) -> InstructionProvenance { diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 16b2c47b5dd..b92161c7329 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -2,16 +2,13 @@ use std::collections::HashSet; use std::sync::Arc; use arc_swap::ArcSwap; +use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecutorFileSystem; -use codex_protocol::error::CodexErr; -use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use futures::FutureExt; -use futures::future::BoxFuture; -use futures::future::Shared; use crate::session::turn_context::TurnEnvironment; use crate::shell::Shell; @@ -31,13 +28,17 @@ pub(crate) fn default_thread_environment_selections( .collect() } -type SnapshotTask = Shared>; +#[derive(Clone, Debug)] +pub(crate) struct StartingTurnEnvironment { + pub(crate) selection: TurnEnvironmentSelection, + pub(crate) environment: Arc, +} pub(crate) struct ThreadEnvironments { environment_manager: Arc, local_shell: Shell, shell_snapshot: ShellSnapshot, - snapshot_task: ArcSwap, + snapshot: ArcSwap, } impl ThreadEnvironments { @@ -51,127 +52,180 @@ impl ThreadEnvironments { environment_manager, local_shell, shell_snapshot, - snapshot_task: ArcSwap::from_pointee(futures::future::ready(current).boxed().shared()), + snapshot: ArcSwap::from_pointee(current), } } pub(crate) fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { - let previous = self - .snapshot_task - .load() - .peek() - .cloned() - .unwrap_or_default(); - let environment_manager = Arc::clone(&self.environment_manager); - let local_shell = self.local_shell.clone(); - let shell_snapshot = self.shell_snapshot.clone(); - let environments = environments.to_vec(); - let (snapshot_task, snapshot) = async move { - Self::resolve_snapshot( - environment_manager, - local_shell, - shell_snapshot, - previous, - environments, - ) - .await - } - .remote_handle(); - self.snapshot_task - .store(Arc::new(snapshot.boxed().shared())); - drop(tokio::spawn(snapshot_task)); - } - - async fn resolve_snapshot( - environment_manager: Arc, - local_shell: Shell, - shell_snapshot: ShellSnapshot, - current: TurnEnvironmentSnapshot, - environments: Vec, - ) -> TurnEnvironmentSnapshot { + let previous = self.snapshot.load(); let mut seen_environment_ids = HashSet::with_capacity(environments.len()); let mut turn_environments = Vec::with_capacity(environments.len()); - for selected_environment in &environments { + let mut starting = Vec::with_capacity(environments.len()); + let mut ordered_selections = Vec::with_capacity(environments.len()); + for selected_environment in environments { if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { continue; } - let turn_environment = match current.turn_environments.iter().find(|environment| { + // Reuse the exact attached or starting environment already selected by this thread. + if let Some(environment) = previous.turn_environments.iter().find(|environment| { environment.environment_id == selected_environment.environment_id && environment.cwd() == &selected_environment.cwd }) { - Some(environment) => environment.clone(), - None => match Self::resolve_selection( - &environment_manager, - &local_shell, - &shell_snapshot, - selected_environment, - ) - .await - { - Ok(environment) => environment, - Err(err) => { - tracing::warn!( - "skipping unresolved turn environment `{}`: {err}", - selected_environment.environment_id - ); - continue; - } - }, + turn_environments.push(environment.clone()); + ordered_selections.push(selected_environment.clone()); + continue; + } + if let Some(environment) = previous + .starting + .iter() + .find(|environment| environment.selection == *selected_environment) + { + starting.push(environment.clone()); + ordered_selections.push(selected_environment.clone()); + continue; + } + + // Only new selections consult the manager; reused selections keep their stable handle. + let environment_id = &selected_environment.environment_id; + let Some(environment) = self.environment_manager.get_environment(environment_id) else { + tracing::warn!("skipping unknown turn environment `{environment_id}`"); + continue; }; - turn_environments.push(turn_environment); + if environment.is_remote() { + // Connect in the background and leave attachment to a later snapshot. + environment.start_connecting(); + starting.push(StartingTurnEnvironment { + selection: selected_environment.clone(), + environment, + }); + } else { + turn_environments.push(self.build_turn_environment( + selected_environment, + environment, + Some(self.local_shell.clone()), + )); + } + ordered_selections.push(selected_environment.clone()); } - TurnEnvironmentSnapshot { turn_environments } + self.snapshot.store(Arc::new(TurnEnvironmentSnapshot { + turn_environments, + starting, + ordered_selections, + })); } - async fn resolve_selection( - environment_manager: &EnvironmentManager, - local_shell: &Shell, - shell_snapshot: &ShellSnapshot, - selected_environment: &TurnEnvironmentSelection, - ) -> CodexResult { - let environment_id = selected_environment.environment_id.clone(); - let environment = environment_manager - .get_environment(&environment_id) - .ok_or_else(|| { - CodexErr::InvalidRequest(format!("unknown turn environment id `{environment_id}`")) - })?; - let shell = if environment.is_remote() { - match environment.info().await { - Ok(info) => match Shell::from_environment_shell_info(info.shell) { - Ok(shell) => Some(shell), - Err(err) => { - tracing::warn!( - "failed to resolve shell for environment `{environment_id}`: {err}" - ); - None - } - }, + async fn resolve_starting_environment( + &self, + starting: &StartingTurnEnvironment, + ) -> TurnEnvironment { + let environment_id = &starting.selection.environment_id; + let shell = match starting.environment.info().boxed().await { + Ok(info) => match Shell::from_environment_shell_info(info.shell) { + Ok(shell) => Some(shell), Err(err) => { - tracing::warn!("failed to get info for environment `{environment_id}`: {err}"); + tracing::warn!( + "failed to resolve shell for environment `{environment_id}`: {err}" + ); None } + }, + Err(err) => { + tracing::warn!("failed to get info for environment `{environment_id}`: {err}"); + None } - } else { - Some(local_shell.clone()) }; + self.build_turn_environment( + &starting.selection, + Arc::clone(&starting.environment), + shell, + ) + } + + fn build_turn_environment( + &self, + selected_environment: &TurnEnvironmentSelection, + environment: Arc, + shell: Option, + ) -> TurnEnvironment { let mut turn_environment = TurnEnvironment::new( - environment_id, + selected_environment.environment_id.clone(), environment, selected_environment.cwd.clone(), shell, ); - let task = shell_snapshot + let task = self + .shell_snapshot .clone() .build(turn_environment.clone()) .boxed() .shared(); drop(tokio::spawn(task.clone())); turn_environment.shell_snapshot = task; - Ok(turn_environment) + turn_environment } pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot { - self.snapshot_task.load_full().as_ref().clone().await + loop { + let current = self.snapshot.load_full(); + if current.starting.is_empty() { + return current.as_ref().clone(); + } + + // Rebuild both lists in configured order while promoting completed startups. + let mut changed = false; + let mut turn_environments = Vec::with_capacity(current.ordered_selections.len()); + let mut starting = Vec::with_capacity(current.starting.len()); + for selection in ¤t.ordered_selections { + if let Some(environment) = current.turn_environments.iter().find(|environment| { + environment.environment_id == selection.environment_id + && environment.cwd() == &selection.cwd + }) { + turn_environments.push(environment.clone()); + continue; + } + let Some(environment) = current + .starting + .iter() + .find(|environment| environment.selection == *selection) + else { + continue; + }; + if !environment.environment.startup_finished() { + // Never wait for an environment whose startup is still running. + starting.push(environment.clone()); + continue; + } + + changed = true; + // Startup finished, so this only reads its saved success or failure. + match environment.environment.wait_until_ready().boxed().await { + Ok(()) => { + turn_environments + .push(self.resolve_starting_environment(environment).await); + } + Err(err) => { + tracing::warn!( + "turn environment `{}` failed to start: {err}", + environment.selection.environment_id + ); + } + } + } + if !changed { + return current.as_ref().clone(); + } + + let next = Arc::new(TurnEnvironmentSnapshot { + turn_environments, + starting, + ordered_selections: current.ordered_selections.clone(), + }); + // Do not overwrite selections changed while shell resolution was in flight. + let previous = self.snapshot.compare_and_swap(¤t, Arc::clone(&next)); + if Arc::ptr_eq(&previous, ¤t) { + return next.as_ref().clone(); + } + } } pub(crate) fn environment_manager(&self) -> Arc { @@ -182,9 +236,25 @@ impl ThreadEnvironments { #[derive(Clone, Debug, Default)] pub(crate) struct TurnEnvironmentSnapshot { pub(crate) turn_environments: Vec, + pub(crate) starting: Vec, + // Attached and starting environments are stored separately, so retain their configured order. + ordered_selections: Vec, } impl TurnEnvironmentSnapshot { + #[cfg(test)] + pub(crate) fn from_turn_environments(turn_environments: Vec) -> Self { + let ordered_selections = turn_environments + .iter() + .map(TurnEnvironment::selection) + .collect(); + Self { + turn_environments, + starting: Vec::new(), + ordered_selections, + } + } + pub(crate) fn primary(&self) -> Option<&TurnEnvironment> { self.turn_environments.first() } @@ -202,10 +272,7 @@ impl TurnEnvironmentSnapshot { } pub(crate) fn to_selections(&self) -> Vec { - self.turn_environments - .iter() - .map(TurnEnvironment::selection) - .collect() + self.ordered_selections.clone() } pub(crate) fn primary_filesystem(&self) -> Option> { @@ -237,7 +304,16 @@ mod tests { use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; + use futures::SinkExt; + use futures::StreamExt; use pretty_assertions::assert_eq; + use serde_json::Value; + use tokio::net::TcpListener; + use tokio::net::TcpStream; + use tokio::time::timeout; + use tokio_tungstenite::WebSocketStream; + use tokio_tungstenite::accept_async; + use tokio_tungstenite::tungstenite::Message; use super::*; @@ -264,6 +340,61 @@ mod tests { .expect("runtime paths") } + async fn read_websocket_json(websocket: &mut WebSocketStream) -> Value { + loop { + match timeout(std::time::Duration::from_secs(5), websocket.next()) + .await + .expect("websocket read should not time out") + .expect("websocket should stay open") + .expect("websocket frame should read") + { + Message::Text(text) => { + return serde_json::from_str(text.as_ref()).expect("valid JSON-RPC message"); + } + Message::Binary(bytes) => { + return serde_json::from_slice(bytes.as_ref()).expect("valid JSON-RPC message"); + } + Message::Ping(_) | Message::Pong(_) => {} + other => panic!("expected JSON-RPC message, got {other:?}"), + } + } + } + + async fn serve_environment_info(listener: TcpListener) { + let (stream, _) = listener.accept().await.expect("connection"); + let mut websocket = accept_async(stream).await.expect("websocket handshake"); + + let initialize = read_websocket_json(&mut websocket).await; + assert_eq!(initialize["method"], "initialize"); + websocket + .send(Message::Text( + serde_json::json!({ + "id": initialize["id"], + "result": { "sessionId": "test-session" } + }) + .to_string() + .into(), + )) + .await + .expect("initialize response"); + let initialized = read_websocket_json(&mut websocket).await; + assert_eq!(initialized["method"], "initialized"); + + let info = read_websocket_json(&mut websocket).await; + assert_eq!(info["method"], "environment/info"); + websocket + .send(Message::Text( + serde_json::json!({ + "id": info["id"], + "result": { "shell": { "name": "zsh", "path": "/bin/zsh" } } + }) + .to_string() + .into(), + )) + .await + .expect("environment info response"); + } + #[tokio::test] async fn default_thread_environment_selections_use_manager_default_id() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); @@ -447,6 +578,84 @@ url = "ws://127.0.0.1:8765" assert_eq!(resolved.snapshot().await.to_selections(), vec![local]); } + #[tokio::test] + async fn snapshot_keeps_starting_environment_until_it_can_be_attached() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind websocket listener"); + let manager = Arc::new( + EnvironmentManager::create_for_tests_with_local( + Some(format!( + "ws://{}", + listener.local_addr().expect("listener address") + )), + test_runtime_paths(), + ) + .await, + ); + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let cwd = PathUri::from_abs_path(&cwd); + let remote = TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: cwd.clone(), + }; + let local = TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd, + }; + let turn_environments = ThreadEnvironments::new( + manager, + crate::shell::default_user_shell(), + ShellSnapshot::disabled(), + TurnEnvironmentSnapshot::default(), + ); + turn_environments.update_selections(&[remote.clone(), local.clone()]); + + let starting = turn_environments.snapshot().await; + assert_eq!( + starting + .turn_environments + .iter() + .map(TurnEnvironment::selection) + .collect::>(), + vec![local.clone()] + ); + assert_eq!( + starting + .starting + .iter() + .map(|environment| environment.selection.clone()) + .collect::>(), + vec![remote.clone()] + ); + assert_eq!( + starting.to_selections(), + vec![remote.clone(), local.clone()] + ); + + let server = tokio::spawn(serve_environment_info(listener)); + timeout( + std::time::Duration::from_secs(5), + starting.starting[0].environment.wait_until_ready(), + ) + .await + .expect("environment startup should finish") + .expect("environment startup should succeed"); + let attached = turn_environments.snapshot().await; + + assert!(attached.starting.is_empty()); + assert_eq!( + attached + .turn_environments + .iter() + .map(TurnEnvironment::selection) + .collect::>(), + vec![remote.clone(), local.clone()] + ); + assert_eq!(attached.to_selections(), vec![remote, local]); + server.await.expect("server task"); + } + #[tokio::test] async fn latest_environment_update_wins_while_previous_resolution_is_pending() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") @@ -495,11 +704,17 @@ url = "ws://127.0.0.1:8765" } #[tokio::test] - async fn matching_environment_id_and_cwd_reuse_resolved_environment() { + async fn matching_environment_id_and_cwd_reuse_starting_environment() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let first_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind first listener"); let manager = Arc::new( EnvironmentManager::create_for_tests( - Some("ws://127.0.0.1:8765".to_string()), + Some(format!( + "ws://{}", + first_listener.local_addr().expect("first listener address") + )), Some(test_runtime_paths()), ) .await, @@ -508,41 +723,58 @@ url = "ws://127.0.0.1:8765" environment_id: REMOTE_ENVIRONMENT_ID.to_string(), cwd: PathUri::from_abs_path(&cwd), }; - let initial = - resolve_turn_environments(Arc::clone(&manager), std::slice::from_ref(&selection)).await; + let environments = ThreadEnvironments::new( + Arc::clone(&manager), + crate::shell::default_user_shell(), + ShellSnapshot::disabled(), + TurnEnvironmentSnapshot::default(), + ); + environments.update_selections(std::slice::from_ref(&selection)); + let initial_snapshot = environments.snapshot().await; + let second_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind second listener"); manager .upsert_environment( REMOTE_ENVIRONMENT_ID.to_string(), - "ws://127.0.0.1:9876".to_string(), + format!( + "ws://{}", + second_listener + .local_addr() + .expect("second listener address") + ), ) .expect("replace environment"); - let initial_snapshot = initial.snapshot().await; - initial.update_selections(std::slice::from_ref(&selection)); - let reused_snapshot = initial.snapshot().await; - initial.update_selections(&[TurnEnvironmentSelection { + environments.update_selections(std::slice::from_ref(&selection)); + let reused_snapshot = environments.snapshot().await; + environments.update_selections(&[TurnEnvironmentSelection { cwd: PathUri::from_abs_path(&cwd.join("changed")), ..selection }]); - let changed_snapshot = initial.snapshot().await; + let changed_snapshot = environments.snapshot().await; assert!(Arc::ptr_eq( &initial_snapshot - .primary() + .starting + .first() .expect("initial environment") .environment, &reused_snapshot - .primary() + .starting + .first() .expect("reused environment") .environment, )); assert!(!Arc::ptr_eq( &reused_snapshot - .primary() + .starting + .first() .expect("reused environment") .environment, &changed_snapshot - .primary() + .starting + .first() .expect("changed environment") .environment, )); @@ -566,25 +798,21 @@ url = "ws://127.0.0.1:8765" Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), ); - let remote = TurnEnvironmentSnapshot { - turn_environments: vec![TurnEnvironment::new( + let remote = TurnEnvironmentSnapshot::from_turn_environments(vec![TurnEnvironment::new( + REMOTE_ENVIRONMENT_ID.to_string(), + remote_environment.clone(), + cwd_uri.clone(), + /*shell*/ None, + )]); + let multiple = TurnEnvironmentSnapshot::from_turn_environments(vec![ + local.primary().expect("local environment").clone(), + TurnEnvironment::new( REMOTE_ENVIRONMENT_ID.to_string(), - remote_environment.clone(), - cwd_uri.clone(), + remote_environment, + cwd_uri, /*shell*/ None, - )], - }; - let multiple = TurnEnvironmentSnapshot { - turn_environments: vec![ - local.primary().expect("local environment").clone(), - TurnEnvironment::new( - REMOTE_ENVIRONMENT_ID.to_string(), - remote_environment, - cwd_uri, - /*shell*/ None, - ), - ], - }; + ), + ]); assert_eq!(local.single_local_environment_cwd(), Some(cwd)); assert_eq!(remote.single_local_environment_cwd(), None); From fea9d3799708bdaa0661df358ed222b8193b8391 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Wed, 17 Jun 2026 21:02:24 -0700 Subject: [PATCH 6/8] core: simplify starting environment resolution --- codex-rs/core/src/environment_selection.rs | 370 +++++++++++---------- 1 file changed, 189 insertions(+), 181 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index b92161c7329..6e7da8f4ab6 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,14 +1,18 @@ use std::collections::HashSet; +use std::fmt; use std::sync::Arc; use arc_swap::ArcSwap; use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; +use codex_exec_server::ExecServerError; use codex_exec_server::ExecutorFileSystem; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use futures::FutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; use crate::session::turn_context::TurnEnvironment; use crate::shell::Shell; @@ -28,17 +32,38 @@ pub(crate) fn default_thread_environment_selections( .collect() } -#[derive(Clone, Debug)] +type TurnEnvironmentResolution = + Shared>>>; + +#[derive(Clone)] +struct SelectedTurnEnvironment { + selection: TurnEnvironmentSelection, + // Starting environments resolve in the background without blocking snapshots. + delayed_start: bool, + resolution: TurnEnvironmentResolution, +} + +#[derive(Clone)] pub(crate) struct StartingTurnEnvironment { pub(crate) selection: TurnEnvironmentSelection, - pub(crate) environment: Arc, + resolution: TurnEnvironmentResolution, +} + +impl fmt::Debug for StartingTurnEnvironment { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("StartingTurnEnvironment") + .field("selection", &self.selection) + .field("resolved", &self.resolution.peek().is_some()) + .finish_non_exhaustive() + } } pub(crate) struct ThreadEnvironments { environment_manager: Arc, local_shell: Shell, shell_snapshot: ShellSnapshot, - snapshot: ArcSwap, + environments: ArcSwap>, } impl ThreadEnvironments { @@ -48,184 +73,141 @@ impl ThreadEnvironments { shell_snapshot: ShellSnapshot, current: TurnEnvironmentSnapshot, ) -> Self { + // Reuse only attached environments from the supplied snapshot; drop starting entries. + let environments = current + .turn_environments + .into_iter() + .map(|environment| { + let selection = environment.selection(); + let resolution: TurnEnvironmentResolution = + futures::future::ready(Ok(environment)).boxed().shared(); + SelectedTurnEnvironment { + selection, + delayed_start: false, + resolution, + } + }) + .collect(); Self { environment_manager, local_shell, shell_snapshot, - snapshot: ArcSwap::from_pointee(current), + environments: ArcSwap::from_pointee(environments), } } pub(crate) fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { - let previous = self.snapshot.load(); + let previous = self.environments.load(); let mut seen_environment_ids = HashSet::with_capacity(environments.len()); - let mut turn_environments = Vec::with_capacity(environments.len()); - let mut starting = Vec::with_capacity(environments.len()); - let mut ordered_selections = Vec::with_capacity(environments.len()); + let mut next = Vec::with_capacity(environments.len()); + let mut new_resolutions = Vec::new(); for selected_environment in environments { if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { continue; } - // Reuse the exact attached or starting environment already selected by this thread. - if let Some(environment) = previous.turn_environments.iter().find(|environment| { - environment.environment_id == selected_environment.environment_id - && environment.cwd() == &selected_environment.cwd - }) { - turn_environments.push(environment.clone()); - ordered_selections.push(selected_environment.clone()); - continue; - } if let Some(environment) = previous - .starting .iter() .find(|environment| environment.selection == *selected_environment) { - starting.push(environment.clone()); - ordered_selections.push(selected_environment.clone()); + next.push(environment.clone()); continue; } - // Only new selections consult the manager; reused selections keep their stable handle. let environment_id = &selected_environment.environment_id; let Some(environment) = self.environment_manager.get_environment(environment_id) else { tracing::warn!("skipping unknown turn environment `{environment_id}`"); continue; }; - if environment.is_remote() { - // Connect in the background and leave attachment to a later snapshot. - environment.start_connecting(); - starting.push(StartingTurnEnvironment { - selection: selected_environment.clone(), - environment, - }); - } else { - turn_environments.push(self.build_turn_environment( - selected_environment, - environment, - Some(self.local_shell.clone()), - )); - } - ordered_selections.push(selected_environment.clone()); + let delayed_start = environment.is_remote() && !environment.startup_finished(); + let resolution = Self::resolve_environment( + selected_environment.clone(), + environment, + self.local_shell.clone(), + self.shell_snapshot.clone(), + ); + new_resolutions.push(resolution.clone()); + next.push(SelectedTurnEnvironment { + selection: selected_environment.clone(), + delayed_start, + resolution, + }); + } + self.environments.store(Arc::new(next)); + for resolution in new_resolutions { + drop(tokio::spawn(resolution)); } - self.snapshot.store(Arc::new(TurnEnvironmentSnapshot { - turn_environments, - starting, - ordered_selections, - })); - } - - async fn resolve_starting_environment( - &self, - starting: &StartingTurnEnvironment, - ) -> TurnEnvironment { - let environment_id = &starting.selection.environment_id; - let shell = match starting.environment.info().boxed().await { - Ok(info) => match Shell::from_environment_shell_info(info.shell) { - Ok(shell) => Some(shell), - Err(err) => { - tracing::warn!( - "failed to resolve shell for environment `{environment_id}`: {err}" - ); - None - } - }, - Err(err) => { - tracing::warn!("failed to get info for environment `{environment_id}`: {err}"); - None - } - }; - self.build_turn_environment( - &starting.selection, - Arc::clone(&starting.environment), - shell, - ) } - fn build_turn_environment( - &self, - selected_environment: &TurnEnvironmentSelection, + fn resolve_environment( + selection: TurnEnvironmentSelection, environment: Arc, - shell: Option, - ) -> TurnEnvironment { - let mut turn_environment = TurnEnvironment::new( - selected_environment.environment_id.clone(), - environment, - selected_environment.cwd.clone(), - shell, - ); - let task = self - .shell_snapshot - .clone() - .build(turn_environment.clone()) - .boxed() - .shared(); - drop(tokio::spawn(task.clone())); - turn_environment.shell_snapshot = task; - turn_environment - } - - pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot { - loop { - let current = self.snapshot.load_full(); - if current.starting.is_empty() { - return current.as_ref().clone(); + local_shell: Shell, + shell_snapshot: ShellSnapshot, + ) -> TurnEnvironmentResolution { + async move { + let environment_id = &selection.environment_id; + if let Err(err) = environment.wait_until_ready().await { + tracing::warn!("turn environment `{environment_id}` failed to start: {err}"); + return Err(Arc::new(err)); } - - // Rebuild both lists in configured order while promoting completed startups. - let mut changed = false; - let mut turn_environments = Vec::with_capacity(current.ordered_selections.len()); - let mut starting = Vec::with_capacity(current.starting.len()); - for selection in ¤t.ordered_selections { - if let Some(environment) = current.turn_environments.iter().find(|environment| { - environment.environment_id == selection.environment_id - && environment.cwd() == &selection.cwd - }) { - turn_environments.push(environment.clone()); - continue; - } - let Some(environment) = current - .starting - .iter() - .find(|environment| environment.selection == *selection) - else { - continue; - }; - if !environment.environment.startup_finished() { - // Never wait for an environment whose startup is still running. - starting.push(environment.clone()); - continue; - } - - changed = true; - // Startup finished, so this only reads its saved success or failure. - match environment.environment.wait_until_ready().boxed().await { - Ok(()) => { - turn_environments - .push(self.resolve_starting_environment(environment).await); - } + let shell = if environment.is_remote() { + match environment.info().await { + Ok(info) => match Shell::from_environment_shell_info(info.shell) { + Ok(shell) => Some(shell), + Err(err) => { + tracing::warn!( + "failed to resolve shell for environment `{environment_id}`: {err}" + ); + None + } + }, Err(err) => { tracing::warn!( - "turn environment `{}` failed to start: {err}", - environment.selection.environment_id + "failed to get info for environment `{environment_id}`: {err}" ); + None } } - } - if !changed { - return current.as_ref().clone(); - } + } else { + Some(local_shell) + }; + let mut turn_environment = + TurnEnvironment::new(selection.environment_id, environment, selection.cwd, shell); + let task = shell_snapshot + .build(turn_environment.clone()) + .boxed() + .shared(); + drop(tokio::spawn(task.clone())); + turn_environment.shell_snapshot = task; + Ok(turn_environment) + } + .boxed() + .shared() + } - let next = Arc::new(TurnEnvironmentSnapshot { - turn_environments, - starting, - ordered_selections: current.ordered_selections.clone(), - }); - // Do not overwrite selections changed while shell resolution was in flight. - let previous = self.snapshot.compare_and_swap(¤t, Arc::clone(&next)); - if Arc::ptr_eq(&previous, ¤t) { - return next.as_ref().clone(); + pub(crate) async fn snapshot(&self) -> TurnEnvironmentSnapshot { + let current = self.environments.load_full(); + let mut turn_environments = Vec::with_capacity(current.len()); + let mut starting = Vec::new(); + for environment in current.iter() { + let resolved = if environment.delayed_start { + environment.resolution.peek().cloned() + } else { + Some(environment.resolution.clone().await) + }; + match resolved { + Some(Ok(turn_environment)) => turn_environments.push(turn_environment), + Some(Err(_)) => {} + None => starting.push(StartingTurnEnvironment { + selection: environment.selection.clone(), + resolution: environment.resolution.clone(), + }), } } + TurnEnvironmentSnapshot { + turn_environments, + starting, + } } pub(crate) fn environment_manager(&self) -> Arc { @@ -237,21 +219,14 @@ impl ThreadEnvironments { pub(crate) struct TurnEnvironmentSnapshot { pub(crate) turn_environments: Vec, pub(crate) starting: Vec, - // Attached and starting environments are stored separately, so retain their configured order. - ordered_selections: Vec, } impl TurnEnvironmentSnapshot { #[cfg(test)] pub(crate) fn from_turn_environments(turn_environments: Vec) -> Self { - let ordered_selections = turn_environments - .iter() - .map(TurnEnvironment::selection) - .collect(); Self { turn_environments, starting: Vec::new(), - ordered_selections, } } @@ -272,7 +247,10 @@ impl TurnEnvironmentSnapshot { } pub(crate) fn to_selections(&self) -> Vec { - self.ordered_selections.clone() + self.turn_environments + .iter() + .map(TurnEnvironment::selection) + .collect() } pub(crate) fn primary_filesystem(&self) -> Option> { @@ -628,19 +606,16 @@ url = "ws://127.0.0.1:8765" .collect::>(), vec![remote.clone()] ); - assert_eq!( - starting.to_selections(), - vec![remote.clone(), local.clone()] - ); + assert_eq!(starting.to_selections(), vec![local.clone()]); let server = tokio::spawn(serve_environment_info(listener)); timeout( std::time::Duration::from_secs(5), - starting.starting[0].environment.wait_until_ready(), + starting.starting[0].resolution.clone(), ) .await - .expect("environment startup should finish") - .expect("environment startup should succeed"); + .expect("environment resolution should finish") + .expect("environment resolution should succeed"); let attached = turn_environments.snapshot().await; assert!(attached.starting.is_empty()); @@ -657,7 +632,7 @@ url = "ws://127.0.0.1:8765" } #[tokio::test] - async fn latest_environment_update_wins_while_previous_resolution_is_pending() { + async fn latest_environment_update_replaces_pending_resolution() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("bind websocket listener"); @@ -704,7 +679,7 @@ url = "ws://127.0.0.1:8765" } #[tokio::test] - async fn matching_environment_id_and_cwd_reuse_starting_environment() { + async fn matching_environment_id_and_cwd_reuse_resolution() { let cwd = AbsolutePathBuf::current_dir().expect("cwd"); let first_listener = TcpListener::bind("127.0.0.1:0") .await @@ -754,29 +729,62 @@ url = "ws://127.0.0.1:8765" }]); let changed_snapshot = environments.snapshot().await; + let initial = initial_snapshot + .starting + .first() + .expect("initial environment"); + let reused = reused_snapshot + .starting + .first() + .expect("reused environment"); + let changed = changed_snapshot + .starting + .first() + .expect("changed environment"); + assert!(initial.resolution.ptr_eq(&reused.resolution)); + assert!(!reused.resolution.ptr_eq(&changed.resolution)); + } + + #[tokio::test] + async fn inherited_environment_reuses_parent_handle() { + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let selection = TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&cwd), + }; + let inherited_environment = Arc::new( + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("inherited environment"), + ); + let inherited = TurnEnvironment::new( + selection.environment_id.clone(), + Arc::clone(&inherited_environment), + selection.cwd.clone(), + /*shell*/ None, + ); + let manager = Arc::new(EnvironmentManager::without_environments()); + manager + .upsert_environment( + REMOTE_ENVIRONMENT_ID.to_string(), + "ws://127.0.0.1:9876".to_string(), + ) + .expect("replacement environment"); + let environments = ThreadEnvironments::new( + manager, + crate::shell::default_user_shell(), + ShellSnapshot::disabled(), + TurnEnvironmentSnapshot::from_turn_environments(vec![inherited]), + ); + + environments.update_selections(std::slice::from_ref(&selection)); + let snapshot = environments.snapshot().await; + assert!(Arc::ptr_eq( - &initial_snapshot - .starting - .first() - .expect("initial environment") - .environment, - &reused_snapshot - .starting - .first() - .expect("reused environment") - .environment, - )); - assert!(!Arc::ptr_eq( - &reused_snapshot - .starting - .first() - .expect("reused environment") - .environment, - &changed_snapshot - .starting - .first() - .expect("changed environment") + &snapshot + .primary() + .expect("inherited environment") .environment, + &inherited_environment, )); } From 00c06a542ecdcef38ca833f5adea1a37cc32f132 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Wed, 17 Jun 2026 21:23:12 -0700 Subject: [PATCH 7/8] core: inline environment resolution startup --- codex-rs/core/src/environment_selection.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 6e7da8f4ab6..ddd4ef4f38b 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -100,7 +100,6 @@ impl ThreadEnvironments { let previous = self.environments.load(); let mut seen_environment_ids = HashSet::with_capacity(environments.len()); let mut next = Vec::with_capacity(environments.len()); - let mut new_resolutions = Vec::new(); for selected_environment in environments { if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) { continue; @@ -118,14 +117,14 @@ impl ThreadEnvironments { tracing::warn!("skipping unknown turn environment `{environment_id}`"); continue; }; - let delayed_start = environment.is_remote() && !environment.startup_finished(); + let delayed_start = !environment.startup_finished(); let resolution = Self::resolve_environment( selected_environment.clone(), environment, self.local_shell.clone(), self.shell_snapshot.clone(), ); - new_resolutions.push(resolution.clone()); + drop(tokio::spawn(resolution.clone())); next.push(SelectedTurnEnvironment { selection: selected_environment.clone(), delayed_start, @@ -133,9 +132,6 @@ impl ThreadEnvironments { }); } self.environments.store(Arc::new(next)); - for resolution in new_resolutions { - drop(tokio::spawn(resolution)); - } } fn resolve_environment( From c065ce6a075daa4e3bba335e7450facdff313635 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Thu, 18 Jun 2026 09:24:43 -0700 Subject: [PATCH 8/8] core: add request-scoped environment context --- codex-rs/core/src/codex_delegate.rs | 31 ++- codex-rs/core/src/codex_delegate_tests.rs | 28 +- codex-rs/core/src/codex_thread.rs | 6 +- codex-rs/core/src/compact.rs | 13 +- codex-rs/core/src/compact_remote.rs | 32 ++- codex-rs/core/src/compact_remote_v2.rs | 11 + codex-rs/core/src/compact_tests.rs | 4 +- .../core/src/context/environment_context.rs | 9 +- codex-rs/core/src/context_manager/updates.rs | 13 +- codex-rs/core/src/guardian/review.rs | 18 ++ codex-rs/core/src/guardian/review_session.rs | 13 +- codex-rs/core/src/guardian/tests.rs | 23 ++ codex-rs/core/src/mcp_openai_file.rs | 48 +++- codex-rs/core/src/mcp_tool_call.rs | 37 +++ codex-rs/core/src/mcp_tool_call_tests.rs | 18 +- codex-rs/core/src/prompt_debug.rs | 16 +- codex-rs/core/src/session/handlers.rs | 24 +- codex-rs/core/src/session/mcp.rs | 11 +- codex-rs/core/src/session/mod.rs | 75 +++++- codex-rs/core/src/session/review.rs | 24 +- codex-rs/core/src/session/step_context.rs | 67 +++++ codex-rs/core/src/session/tests.rs | 249 ++++++++++-------- .../core/src/session/tests/guardian_tests.rs | 19 +- codex-rs/core/src/session/turn.rs | 73 +++-- codex-rs/core/src/session/turn_context.rs | 26 +- codex-rs/core/src/session_startup_prewarm.rs | 2 + .../core/src/stream_events_utils_tests.rs | 5 +- codex-rs/core/src/tasks/compact.rs | 7 +- codex-rs/core/src/tasks/review.rs | 14 +- codex-rs/core/src/tasks/user_shell.rs | 4 +- codex-rs/core/src/thread_manager_tests.rs | 32 ++- .../src/thread_rollout_truncation_tests.rs | 2 +- codex-rs/core/src/tools/code_mode/delegate.rs | 143 +++++----- .../src/tools/code_mode/execute_handler.rs | 3 +- codex-rs/core/src/tools/code_mode/mod.rs | 22 +- codex-rs/core/src/tools/context.rs | 2 + .../core/src/tools/handlers/agent_jobs.rs | 5 +- .../agent_jobs/spawn_agents_on_csv.rs | 17 +- .../core/src/tools/handlers/apply_patch.rs | 7 +- .../src/tools/handlers/apply_patch_tests.rs | 3 + .../src/tools/handlers/extension_tools.rs | 14 +- codex-rs/core/src/tools/handlers/mcp.rs | 11 + codex-rs/core/src/tools/handlers/mod.rs | 8 +- .../src/tools/handlers/multi_agents/spawn.rs | 3 +- .../src/tools/handlers/multi_agents_tests.rs | 3 + .../tools/handlers/multi_agents_v2/spawn.rs | 3 +- .../src/tools/handlers/request_permissions.rs | 4 +- .../handlers/request_user_input_tests.rs | 4 + codex-rs/core/src/tools/handlers/shell.rs | 7 +- .../src/tools/handlers/shell/shell_command.rs | 15 +- .../core/src/tools/handlers/shell_tests.rs | 16 +- .../handlers/unified_exec/exec_command.rs | 11 +- .../src/tools/handlers/unified_exec_tests.rs | 11 + .../core/src/tools/handlers/view_image.rs | 20 +- codex-rs/core/src/tools/network_approval.rs | 4 + codex-rs/core/src/tools/orchestrator.rs | 4 + codex-rs/core/src/tools/parallel.rs | 5 + codex-rs/core/src/tools/registry_tests.rs | 3 + codex-rs/core/src/tools/router.rs | 21 +- codex-rs/core/src/tools/router_tests.rs | 24 +- .../core/src/tools/runtimes/apply_patch.rs | 11 +- codex-rs/core/src/tools/runtimes/shell.rs | 1 + .../tools/runtimes/shell/unix_escalation.rs | 5 + .../runtimes/shell/unix_escalation_tests.rs | 8 + .../core/src/tools/runtimes/unified_exec.rs | 1 + codex-rs/core/src/tools/sandboxing.rs | 3 + codex-rs/core/src/tools/spec_plan.rs | 28 +- codex-rs/core/src/tools/spec_plan_tests.rs | 133 ++++++---- .../src/tools/tool_dispatch_trace_tests.rs | 3 + codex-rs/core/src/unified_exec/mod.rs | 10 +- codex-rs/core/src/unified_exec/mod_tests.rs | 19 +- .../core/src/unified_exec/process_manager.rs | 1 + .../src/unified_exec/process_manager_tests.rs | 4 +- codex-rs/core/tests/suite/code_mode.rs | 47 +++- codex-rs/core/tests/suite/tools.rs | 120 +++++++++ 75 files changed, 1288 insertions(+), 453 deletions(-) create mode 100644 codex-rs/core/src/session/step_context.rs diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index b3f0d6c43cd..c9dd964e15f 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -32,6 +32,7 @@ use tokio::time::timeout; use tokio_util::sync::CancellationToken; use crate::config::Config; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::guardian::GuardianApprovalRequest; use crate::guardian::new_guardian_review_id; use crate::guardian::routes_approval_to_guardian; @@ -72,6 +73,7 @@ pub(crate) async fn run_codex_thread_interactive( models_manager: SharedModelsManager, parent_session: Arc, parent_ctx: Arc, + parent_environments: TurnEnvironmentSnapshot, cancel_token: CancellationToken, subagent_source: SubAgentSource, initial_history: Option, @@ -84,6 +86,11 @@ pub(crate) async fn run_codex_thread_interactive( instructions: parent_session.user_instructions().await, warnings: Vec::new(), }; + // Child threads inherit only environments attached for this request. + let attached_environments = TurnEnvironmentSnapshot { + turn_environments: parent_environments.turn_environments.clone(), + starting: Vec::new(), + }; let CodexSpawnOk { codex, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { config, user_instructions, @@ -107,11 +114,11 @@ pub(crate) async fn run_codex_thread_interactive( dynamic_tools: Vec::new(), metrics_service_name: None, user_shell_override: None, - inherited_environments: Some(parent_ctx.environments.clone()), + inherited_environments: Some(attached_environments.clone()), inherited_exec_policy: Some(Arc::clone(&parent_session.services.exec_policy)), parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), parent_trace: None, - environment_selections: parent_ctx.environments.to_selections(), + environment_selections: attached_environments.to_selections(), thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), @@ -141,6 +148,7 @@ pub(crate) async fn run_codex_thread_interactive( // routing them to the parent session for decisions. let parent_session_clone = Arc::clone(&parent_session); let parent_ctx_clone = Arc::clone(&parent_ctx); + let parent_environments_clone = parent_environments.clone(); let codex_for_events = Arc::clone(&codex); // Cache delegated MCP invocations so guardian can recover the full tool call // context when the later legacy RequestUserInput approval event only carries @@ -152,6 +160,7 @@ pub(crate) async fn run_codex_thread_interactive( tx_sub, parent_session_clone, parent_ctx_clone, + parent_environments_clone, pending_mcp_invocations, cancel_token_events, ) @@ -184,6 +193,7 @@ pub(crate) async fn run_codex_thread_one_shot( input: Vec, parent_session: Arc, parent_ctx: Arc, + parent_environments: TurnEnvironmentSnapshot, cancel_token: CancellationToken, subagent_source: SubAgentSource, final_output_json_schema: Option, @@ -198,6 +208,7 @@ pub(crate) async fn run_codex_thread_one_shot( models_manager, parent_session, parent_ctx, + parent_environments, child_cancel.clone(), subagent_source, initial_history, @@ -263,6 +274,7 @@ async fn forward_events( tx_sub: Sender, parent_session: Arc, parent_ctx: Arc, + parent_environments: TurnEnvironmentSnapshot, pending_mcp_invocations: Arc>>, cancel_token: CancellationToken, ) { @@ -299,6 +311,7 @@ async fn forward_events( id, &parent_session, &parent_ctx, + &parent_environments, event, &cancel_token, ) @@ -313,6 +326,7 @@ async fn forward_events( id, &parent_session, &parent_ctx, + &parent_environments, event, &cancel_token, ) @@ -326,6 +340,7 @@ async fn forward_events( &codex, &parent_session, &parent_ctx, + &parent_environments, event, &cancel_token, ) @@ -340,6 +355,7 @@ async fn forward_events( id, &parent_session, &parent_ctx, + &parent_environments, &pending_mcp_invocations, event, &cancel_token, @@ -453,6 +469,7 @@ async fn handle_exec_approval( turn_id: String, parent_session: &Arc, parent_ctx: &Arc, + parent_environments: &TurnEnvironmentSnapshot, event: ExecApprovalRequestEvent, cancel_token: &CancellationToken, ) { @@ -475,6 +492,7 @@ async fn handle_exec_approval( let review_rx = spawn_approval_request_review( Arc::clone(parent_session), Arc::clone(parent_ctx), + parent_environments.clone(), new_guardian_review_id(), GuardianApprovalRequest::Shell { id: call_id.clone(), @@ -538,6 +556,7 @@ async fn handle_patch_approval( _id: String, parent_session: &Arc, parent_ctx: &Arc, + parent_environments: &TurnEnvironmentSnapshot, event: ApplyPatchApprovalRequestEvent, cancel_token: &CancellationToken, ) { @@ -588,6 +607,7 @@ async fn handle_patch_approval( let review_rx = spawn_approval_request_review( Arc::clone(parent_session), Arc::clone(parent_ctx), + parent_environments.clone(), new_guardian_review_id(), GuardianApprovalRequest::ApplyPatch { id: approval_id.clone(), @@ -636,11 +656,13 @@ async fn handle_patch_approval( .await; } +#[allow(clippy::too_many_arguments)] async fn handle_request_user_input( codex: &Codex, id: String, parent_session: &Arc, parent_ctx: &Arc, + parent_environments: &TurnEnvironmentSnapshot, pending_mcp_invocations: &Arc>>, event: RequestUserInputEvent, cancel_token: &CancellationToken, @@ -648,6 +670,7 @@ async fn handle_request_user_input( if let Some(response) = maybe_auto_review_mcp_request_user_input( parent_session, parent_ctx, + parent_environments, pending_mcp_invocations, &event, cancel_token, @@ -684,6 +707,7 @@ async fn handle_request_user_input( async fn maybe_auto_review_mcp_request_user_input( parent_session: &Arc, parent_ctx: &Arc, + parent_environments: &TurnEnvironmentSnapshot, pending_mcp_invocations: &Arc>>, event: &RequestUserInputEvent, cancel_token: &CancellationToken, @@ -716,6 +740,7 @@ async fn maybe_auto_review_mcp_request_user_input( let review_rx = spawn_approval_request_review( Arc::clone(parent_session), Arc::clone(parent_ctx), + parent_environments.clone(), new_guardian_review_id(), build_guardian_mcp_tool_review_request(&event.call_id, &invocation, metadata.as_ref()), /*retry_reason*/ None, @@ -762,6 +787,7 @@ async fn handle_request_permissions( codex: &Codex, parent_session: &Arc, parent_ctx: &Arc, + parent_environments: &TurnEnvironmentSnapshot, event: RequestPermissionsEvent, cancel_token: &CancellationToken, ) { @@ -777,6 +803,7 @@ async fn handle_request_permissions( }); let response_fut = parent_session.request_permissions_for_cwd( parent_ctx, + parent_environments, call_id.clone(), args, cwd, diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 21a86f224e0..d59f5a477f5 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -39,6 +39,7 @@ async fn forward_events_cancelled_while_send_blocked_shuts_down_delegate() { let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); let (session, ctx, _rx_evt) = crate::session::tests::make_session_and_context_with_rx().await; + let step = crate::session::tests::step_context_for_session(session.as_ref()).await; let codex = Arc::new(Codex { tx_sub, rx_event: rx_events, @@ -67,6 +68,7 @@ async fn forward_events_cancelled_while_send_blocked_shuts_down_delegate() { tx_out.clone(), session, ctx, + step.environments.clone(), Arc::new(Mutex::new(HashMap::new())), cancel.clone(), )); @@ -160,6 +162,8 @@ async fn forward_ops_preserves_submission_trace_context() { async fn run_codex_thread_interactive_respects_pre_cancelled_spawn() { let (parent_session, parent_ctx, _rx_events) = crate::session::tests::make_session_and_context_with_rx().await; + let parent_step = + crate::session::tests::step_context_for_session(parent_session.as_ref()).await; let cancel_token = CancellationToken::new(); cancel_token.cancel(); @@ -171,6 +175,7 @@ async fn run_codex_thread_interactive_respects_pre_cancelled_spawn() { Arc::clone(&parent_session.services.models_manager), parent_session, parent_ctx, + parent_step.environments.clone(), cancel_token, SubAgentSource::Review, /*initial_history*/ None, @@ -184,11 +189,16 @@ async fn run_codex_thread_interactive_respects_pre_cancelled_spawn() { #[tokio::test] async fn handle_request_permissions_uses_tool_call_id_for_round_trip() { - let (parent_session, mut parent_ctx, rx_events) = + let (parent_session, parent_ctx, rx_events) = crate::session::tests::make_session_and_context_with_rx().await; + let mut parent_step = + crate::session::tests::step_context_for_session(parent_session.as_ref()).await; *parent_session.active_turn.lock().await = Some(crate::state::ActiveTurn::default()); - let parent_ctx_mut = Arc::get_mut(&mut parent_ctx).expect("single turn context ref"); - parent_ctx_mut.environments.turn_environments[0].environment_id = "remote".to_string(); + Arc::get_mut(&mut parent_step) + .expect("single step context ref") + .environments + .turn_environments[0] + .environment_id = "remote".to_string(); let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_tx_events, rx_events_child) = bounded(SUBMISSION_CHANNEL_CAPACITY); @@ -222,12 +232,14 @@ async fn handle_request_permissions_uses_tool_call_id_for_round_trip() { let codex = Arc::clone(&codex); let parent_session = Arc::clone(&parent_session); let parent_ctx = Arc::clone(&parent_ctx); + let parent_step = Arc::clone(&parent_step); let cancel_token = cancel_token.clone(); async move { handle_request_permissions( codex.as_ref(), &parent_session, &parent_ctx, + &parent_step.environments, RequestPermissionsEvent { call_id: request_call_id, turn_id: "child-turn-1".to_string(), @@ -307,10 +319,13 @@ async fn handle_exec_approval_uses_call_id_for_guardian_review_and_approval_id_f }); let cancel_token = CancellationToken::new(); + let parent_step = + crate::session::tests::step_context_for_session(parent_session.as_ref()).await; let handle = tokio::spawn({ let codex = Arc::clone(&codex); let parent_session = Arc::clone(&parent_session); let parent_ctx = Arc::clone(&parent_ctx); + let parent_step = Arc::clone(&parent_step); let cancel_token = cancel_token.clone(); async move { handle_exec_approval( @@ -318,6 +333,7 @@ async fn handle_exec_approval_uses_call_id_for_guardian_review_and_approval_id_f "child-turn-1".to_string(), &parent_session, &parent_ctx, + &parent_step.environments, ExecApprovalRequestEvent { call_id: "command-item-1".to_string(), approval_id: Some("callback-approval-1".to_string()), @@ -419,10 +435,13 @@ async fn delegated_mcp_guardian_abort_returns_synthetic_decline_answer() { )]))); let cancel_token = CancellationToken::new(); cancel_token.cancel(); + let parent_step = + crate::session::tests::step_context_for_session(parent_session.as_ref()).await; let response = maybe_auto_review_mcp_request_user_input( &parent_session, &parent_ctx, + &parent_step.environments, &pending_mcp_invocations, &RequestUserInputEvent { call_id: "call-1".to_string(), @@ -467,6 +486,8 @@ async fn delegated_mcp_user_reviewer_returns_none_without_metadata() { }, )]))); let cancel_token = CancellationToken::new(); + let parent_step = + crate::session::tests::step_context_for_session(parent_session.as_ref()).await; let event = RequestUserInputEvent { call_id: "call-1".to_string(), @@ -484,6 +505,7 @@ async fn delegated_mcp_user_reviewer_returns_none_without_metadata() { let response = maybe_auto_review_mcp_request_user_input( &parent_session, &parent_ctx, + &parent_step.environments, &pending_mcp_invocations, &event, &cancel_token, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 99e0391d3b0..6b331acd511 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -456,9 +456,13 @@ impl CodexThread { let turn_context = self.codex.session.new_default_turn().await; if self.codex.session.reference_context_item().await.is_none() { + let step_context = self.codex.session.prepare_step_for_request().await; self.codex .session - .record_context_updates_and_set_reference_context_item(turn_context.as_ref()) + .record_context_updates_and_set_reference_context_item( + turn_context.as_ref(), + step_context.as_ref(), + ) .await; } self.codex diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 9d79138e87c..6d011877649 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -14,6 +14,7 @@ use crate::responses_metadata::CompactionTurnMetadata; #[cfg(test)] use crate::session::PreviousTurnSettings; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn::get_last_assistant_message_from_turn; use crate::session::turn_context::TurnContext; use crate::util::backoff; @@ -73,6 +74,7 @@ pub(crate) fn should_use_remote_compact_task(provider: &ModelProviderInfo) -> bo pub(crate) async fn run_inline_auto_compact_task( sess: Arc, turn_context: Arc, + step_context: Arc, initial_context_injection: InitialContextInjection, reason: CompactionReason, phase: CompactionPhase, @@ -92,6 +94,7 @@ pub(crate) async fn run_inline_auto_compact_task( run_compact_task_inner( sess, turn_context, + step_context, input, initial_context_injection, CompactionTrigger::Auto, @@ -105,6 +108,7 @@ pub(crate) async fn run_inline_auto_compact_task( pub(crate) async fn run_compact_task( sess: Arc, turn_context: Arc, + step_context: Arc, input: Vec, ) -> CodexResult<()> { let start_event = EventMsg::TurnStarted(TurnStartedEvent { @@ -118,6 +122,7 @@ pub(crate) async fn run_compact_task( run_compact_task_inner( sess.clone(), turn_context, + step_context, input, InitialContextInjection::DoNotInject, CompactionTrigger::Manual, @@ -128,9 +133,11 @@ pub(crate) async fn run_compact_task( Ok(()) } +#[allow(clippy::too_many_arguments)] async fn run_compact_task_inner( sess: Arc, turn_context: Arc, + step_context: Arc, input: Vec, initial_context_injection: InitialContextInjection, trigger: CompactionTrigger, @@ -167,6 +174,7 @@ async fn run_compact_task_inner( let result = run_compact_task_inner_impl( Arc::clone(&sess), Arc::clone(&turn_context), + Arc::clone(&step_context), input, initial_context_injection, compaction_metadata, @@ -202,6 +210,7 @@ async fn run_compact_task_inner( async fn run_compact_task_inner_impl( sess: Arc, turn_context: Arc, + step_context: Arc, input: Vec, initial_context_injection: InitialContextInjection, compaction_metadata: CompactionTurnMetadata, @@ -308,7 +317,9 @@ async fn run_compact_task_inner_impl( initial_context_injection, InitialContextInjection::BeforeLastUserMessage ) { - let initial_context = sess.build_initial_context(turn_context.as_ref()).await; + let initial_context = sess + .build_initial_context(turn_context.as_ref(), step_context.as_ref()) + .await; new_history = insert_initial_context_before_last_real_user_or_summary(new_history, initial_context); } diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 08ca7efcdfe..7a3db82539d 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -16,6 +16,7 @@ use crate::hook_runtime::run_pre_compact_hooks; use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_metadata::CompactionTurnMetadata; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn::built_tools; use crate::session::turn_context::TurnContext; use codex_analytics::CompactionImplementation; @@ -44,6 +45,7 @@ const CONTEXT_WINDOW_TRUNCATED_OUTPUT_MESSAGE: &str = pub(crate) async fn run_inline_remote_auto_compact_task( sess: Arc, turn_context: Arc, + step_context: Arc, turn_state: Arc>, initial_context_injection: InitialContextInjection, reason: CompactionReason, @@ -52,6 +54,7 @@ pub(crate) async fn run_inline_remote_auto_compact_task( run_remote_compact_task_inner( &sess, &turn_context, + &step_context, Some(turn_state), initial_context_injection, CompactionTrigger::Auto, @@ -65,6 +68,7 @@ pub(crate) async fn run_inline_remote_auto_compact_task( pub(crate) async fn run_remote_compact_task( sess: Arc, turn_context: Arc, + step_context: Arc, ) -> CodexResult<()> { let start_event = EventMsg::TurnStarted(TurnStartedEvent { turn_id: turn_context.sub_id.clone(), @@ -78,6 +82,7 @@ pub(crate) async fn run_remote_compact_task( run_remote_compact_task_inner( &sess, &turn_context, + &step_context, /*turn_state*/ None, InitialContextInjection::DoNotInject, CompactionTrigger::Manual, @@ -88,9 +93,11 @@ pub(crate) async fn run_remote_compact_task( Ok(()) } +#[allow(clippy::too_many_arguments)] async fn run_remote_compact_task_inner( sess: &Arc, turn_context: &Arc, + step_context: &Arc, turn_state: Option>>, initial_context_injection: InitialContextInjection, trigger: CompactionTrigger, @@ -135,6 +142,7 @@ async fn run_remote_compact_task_inner( let result = run_remote_compact_task_inner_impl( sess, turn_context, + step_context, turn_state, initial_context_injection, compaction_metadata, @@ -169,6 +177,7 @@ async fn run_remote_compact_task_inner( async fn run_remote_compact_task_inner_impl( sess: &Arc, turn_context: &Arc, + step_context: &Arc, turn_state: Option>>, initial_context_injection: InitialContextInjection, compaction_metadata: CompactionTurnMetadata, @@ -220,6 +229,7 @@ async fn run_remote_compact_task_inner_impl( let tool_router = built_tools( sess.as_ref(), turn_context.as_ref(), + Arc::clone(step_context), &CancellationToken::new(), ) .await?; @@ -262,6 +272,7 @@ async fn run_remote_compact_task_inner_impl( new_history = process_compacted_history( sess.as_ref(), turn_context.as_ref(), + step_context.as_ref(), new_history, initial_context_injection, ) @@ -295,6 +306,7 @@ async fn run_remote_compact_task_inner_impl( pub(crate) async fn process_compacted_history( sess: &Session, turn_context: &TurnContext, + step_context: &StepContext, mut compacted_history: Vec, initial_context_injection: InitialContextInjection, ) -> Vec { @@ -305,7 +317,7 @@ pub(crate) async fn process_compacted_history( initial_context_injection, InitialContextInjection::BeforeLastUserMessage ) { - sess.build_initial_context(turn_context).await + sess.build_initial_context(turn_context, step_context).await } else { Vec::new() }; @@ -314,6 +326,24 @@ pub(crate) async fn process_compacted_history( insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context) } +#[cfg(test)] +pub(crate) async fn process_compacted_history_for_test( + sess: &Session, + turn_context: &TurnContext, + compacted_history: Vec, + initial_context_injection: InitialContextInjection, +) -> Vec { + let step_context = StepContext::new(sess.services.turn_environments.snapshot().await); + process_compacted_history( + sess, + turn_context, + &step_context, + compacted_history, + initial_context_injection, + ) + .await +} + /// Returns whether an item from remote compaction output should be preserved. /// /// Called while processing the model-provided compacted transcript, before we diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 4b4071a0d97..eadd461c000 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -21,6 +21,7 @@ use crate::responses_metadata::CompactionTurnMetadata; use crate::responses_retry::ResponsesStreamRequest; use crate::responses_retry::handle_retryable_response_stream_error; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn::built_tools; use crate::session::turn_context::TurnContext; use codex_analytics::CompactionImplementation; @@ -56,6 +57,7 @@ const MAX_REMOTE_COMPACTION_V2_STREAM_RETRIES: u64 = 2; pub(crate) async fn run_inline_remote_auto_compact_task( sess: Arc, turn_context: Arc, + step_context: Arc, client_session: &mut ModelClientSession, initial_context_injection: InitialContextInjection, reason: CompactionReason, @@ -64,6 +66,7 @@ pub(crate) async fn run_inline_remote_auto_compact_task( run_remote_compact_task_inner( &sess, &turn_context, + &step_context, Some(client_session), initial_context_injection, CompactionTrigger::Auto, @@ -76,6 +79,7 @@ pub(crate) async fn run_inline_remote_auto_compact_task( pub(crate) async fn run_remote_compact_task( sess: Arc, turn_context: Arc, + step_context: Arc, ) -> CodexResult<()> { let start_event = EventMsg::TurnStarted(TurnStartedEvent { turn_id: turn_context.sub_id.clone(), @@ -89,6 +93,7 @@ pub(crate) async fn run_remote_compact_task( run_remote_compact_task_inner( &sess, &turn_context, + &step_context, /*client_session*/ None, InitialContextInjection::DoNotInject, CompactionTrigger::Manual, @@ -98,9 +103,11 @@ pub(crate) async fn run_remote_compact_task( .await } +#[allow(clippy::too_many_arguments)] async fn run_remote_compact_task_inner( sess: &Arc, turn_context: &Arc, + step_context: &Arc, client_session: Option<&mut ModelClientSession>, initial_context_injection: InitialContextInjection, trigger: CompactionTrigger, @@ -145,6 +152,7 @@ async fn run_remote_compact_task_inner( let result = run_remote_compact_task_inner_impl( sess, turn_context, + step_context, client_session, initial_context_injection, compaction_metadata, @@ -179,6 +187,7 @@ async fn run_remote_compact_task_inner( async fn run_remote_compact_task_inner_impl( sess: &Arc, turn_context: &Arc, + step_context: &Arc, client_session: Option<&mut ModelClientSession>, initial_context_injection: InitialContextInjection, compaction_metadata: CompactionTurnMetadata, @@ -227,6 +236,7 @@ async fn run_remote_compact_task_inner_impl( let tool_router = built_tools( sess.as_ref(), turn_context.as_ref(), + Arc::clone(step_context), &CancellationToken::new(), ) .await?; @@ -292,6 +302,7 @@ async fn run_remote_compact_task_inner_impl( let new_history = process_compacted_history( sess.as_ref(), turn_context.as_ref(), + step_context.as_ref(), compacted_history, initial_context_injection, ) diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index c1004004ade..7ec440d445a 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -13,8 +13,8 @@ async fn process_compacted_history_with_test_session( session .set_previous_turn_settings(previous_turn_settings.cloned()) .await; - let initial_context = session.build_initial_context(&turn_context).await; - let refreshed = crate::compact_remote::process_compacted_history( + let initial_context = session.build_initial_context_for_test(&turn_context).await; + let refreshed = crate::compact_remote::process_compacted_history_for_test( &session, &turn_context, compacted_history, diff --git a/codex-rs/core/src/context/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index e93ffb88ffd..08b90900204 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -1,3 +1,4 @@ +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::session::turn_context::TurnEnvironment; use crate::shell::Shell; @@ -419,10 +420,14 @@ impl EnvironmentContext { ) } - pub(crate) fn from_turn_context(turn_context: &TurnContext, shell: &Shell) -> Self { + pub(crate) fn from_turn_context( + turn_context: &TurnContext, + step_context: &StepContext, + shell: &Shell, + ) -> Self { let mut context = Self::new( EnvironmentContextEnvironment::from_turn_environments( - &turn_context.environments.turn_environments, + &step_context.environments.turn_environments, shell, ), turn_context.current_date.clone(), diff --git a/codex-rs/core/src/context_manager/updates.rs b/codex-rs/core/src/context_manager/updates.rs index 4c69aae95b2..58b39b32527 100644 --- a/codex-rs/core/src/context_manager/updates.rs +++ b/codex-rs/core/src/context_manager/updates.rs @@ -8,6 +8,7 @@ use crate::context::RealtimeEndInstructions; use crate::context::RealtimeStartInstructions; use crate::context::RealtimeStartWithInstructions; use crate::session::PreviousTurnSettings; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::shell::Shell; use codex_execpolicy::Policy; @@ -21,6 +22,7 @@ use codex_protocol::protocol::TurnContextItem; fn build_environment_update_item( previous: Option<&TurnContextItem>, next: &TurnContext, + step: &StepContext, shell: &Shell, ) -> Option { if !next.config.include_environment_context { @@ -29,7 +31,7 @@ fn build_environment_update_item( let prev = previous?; let prev_context = EnvironmentContext::from_turn_context_item(prev, shell.name().to_string()); - let next_context = EnvironmentContext::from_turn_context(next, shell); + let next_context = EnvironmentContext::from_turn_context(next, step, shell); if prev_context.equals_except_shell(&next_context) { return None; } @@ -42,6 +44,7 @@ fn build_environment_update_item( fn build_permissions_update_item( previous: Option<&TurnContextItem>, next: &TurnContext, + step: &StepContext, exec_policy: &Policy, ) -> Option { if !next.config.include_permissions_instructions { @@ -61,8 +64,7 @@ fn build_permissions_update_item( next.approval_policy.value(), next.config.approvals_reviewer, exec_policy, - #[allow(deprecated)] - &next.cwd, + &step.effective_cwd(next), next.config .features .enabled(Feature::ExecPermissionApprovals), @@ -215,6 +217,7 @@ pub(crate) fn build_settings_update_items( previous: Option<&TurnContextItem>, previous_turn_settings: Option<&PreviousTurnSettings>, next: &TurnContext, + step: &StepContext, shell: &Shell, exec_policy: &Policy, personality_feature_enabled: bool, @@ -223,12 +226,12 @@ pub(crate) fn build_settings_update_items( // model-visible item emitted by build_initial_context. Persist the remaining // inputs or add explicit replay events so fork/resume can diff everything // deterministically. - let contextual_user_message = build_environment_update_item(previous, next, shell); + let contextual_user_message = build_environment_update_item(previous, next, step, shell); let developer_update_sections = [ // Keep model-switch instructions first so model-specific guidance is read before // any other context diffs on this turn. build_model_instructions_update_item(previous_turn_settings, next), - build_permissions_update_item(previous, next, exec_policy), + build_permissions_update_item(previous, next, step, exec_policy), build_collaboration_mode_update_item(previous, next), build_realtime_update_item(previous, previous_turn_settings, next), build_personality_update_item(previous, next, personality_feature_enabled), diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 92223083d8d..3808ef208fd 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -24,6 +24,7 @@ use tokio::time::Instant; use tokio::time::sleep_until; use tokio_util::sync::CancellationToken; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::turn_timing::now_unix_timestamp_ms; @@ -272,9 +273,11 @@ pub(crate) async fn record_guardian_denial_for_test( /// This function always fails closed: timeouts, review-session failures, and /// parse failures all block execution, but timeouts are still surfaced to the /// caller as distinct from explicit guardian denials. +#[allow(clippy::too_many_arguments)] async fn run_guardian_review( session: Arc, turn: Arc, + environments: TurnEnvironmentSnapshot, review_id: String, request: GuardianApprovalRequest, retry_reason: Option, @@ -359,6 +362,7 @@ async fn run_guardian_review( let (outcome, analytics_result) = Box::pin(run_guardian_review_session_with_retry( session.clone(), turn.clone(), + environments, request, retry_reason.clone(), schema, @@ -594,6 +598,7 @@ async fn run_guardian_review( pub(crate) async fn review_approval_request( session: &Arc, turn: &Arc, + environments: TurnEnvironmentSnapshot, review_id: String, request: GuardianApprovalRequest, retry_reason: Option, @@ -603,6 +608,7 @@ pub(crate) async fn review_approval_request( Box::pin(run_guardian_review( Arc::clone(session), Arc::clone(turn), + environments, review_id, request, retry_reason, @@ -612,9 +618,11 @@ pub(crate) async fn review_approval_request( .await } +#[allow(clippy::too_many_arguments)] pub(crate) async fn review_approval_request_with_cancel( session: &Arc, turn: &Arc, + environments: TurnEnvironmentSnapshot, review_id: String, request: GuardianApprovalRequest, retry_reason: Option, @@ -624,6 +632,7 @@ pub(crate) async fn review_approval_request_with_cancel( run_guardian_review( Arc::clone(session), Arc::clone(turn), + environments, review_id, request, retry_reason, @@ -633,9 +642,11 @@ pub(crate) async fn review_approval_request_with_cancel( .await } +#[allow(clippy::too_many_arguments)] pub(crate) fn spawn_approval_request_review( session: Arc, turn: Arc, + environments: TurnEnvironmentSnapshot, review_id: String, request: GuardianApprovalRequest, retry_reason: Option, @@ -654,6 +665,7 @@ pub(crate) fn spawn_approval_request_review( let decision = runtime.block_on(review_approval_request_with_cancel( &session, &turn, + environments, review_id, request, retry_reason, @@ -679,9 +691,11 @@ pub(crate) fn spawn_approval_request_review( /// context. It may still reuse the parent's managed-network allowlist for /// read-only checks, but it intentionally runs without inherited exec-policy /// rules. +#[allow(clippy::too_many_arguments)] async fn run_guardian_review_session_before_deadline( session: Arc, turn: Arc, + environments: TurnEnvironmentSnapshot, request: GuardianApprovalRequest, retry_reason: Option, schema: serde_json::Value, @@ -772,6 +786,7 @@ async fn run_guardian_review_session_before_deadline( .run_review(GuardianReviewSessionParams { parent_session: Arc::clone(&session), parent_turn: turn.clone(), + parent_environments: environments, spawn_config: guardian_config, request, retry_reason, @@ -841,9 +856,11 @@ async fn run_guardian_review_session_before_deadline( } } +#[allow(clippy::too_many_arguments)] pub(super) async fn run_guardian_review_session_with_retry( session: Arc, turn: Arc, + environments: TurnEnvironmentSnapshot, request: GuardianApprovalRequest, retry_reason: Option, schema: serde_json::Value, @@ -857,6 +874,7 @@ pub(super) async fn run_guardian_review_session_with_retry( let (outcome, mut analytics_result) = run_guardian_review_session_before_deadline( Arc::clone(&session), Arc::clone(&turn), + environments.clone(), request.clone(), retry_reason.clone(), schema.clone(), diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 7aa007891f2..537fddc323d 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -41,6 +41,7 @@ use crate::config::NetworkProxySpec; use crate::config::Permissions; use crate::context::ContextualUserFragment; use crate::context::GuardianFollowupReviewReminder; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::session::Codex; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -73,6 +74,7 @@ pub(crate) enum GuardianReviewSessionOutcome { pub(crate) struct GuardianReviewSessionParams { pub(crate) parent_session: Arc, pub(crate) parent_turn: Arc, + pub(crate) parent_environments: TurnEnvironmentSnapshot, pub(crate) spawn_config: Config, pub(crate) request: GuardianApprovalRequest, pub(crate) retry_reason: Option, @@ -625,6 +627,7 @@ async fn spawn_guardian_review_session( params.parent_session.services.models_manager.clone(), Arc::clone(¶ms.parent_session), Arc::clone(¶ms.parent_turn), + params.parent_environments.clone(), cancel_token.clone(), SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_string()), initial_history, @@ -749,12 +752,11 @@ async fn run_review_on_session( .await .unwrap_or_default(); let guardian_permission_profile = PermissionProfile::read_only(); - let parent_turn_environments = params.parent_turn.environments.to_selections(); + let parent_turn_environments = params.parent_environments.to_selections(); // TODO(anp): Migrate guardian review thread settings to a PathUri fallback cwd so foreign // parent environments do not fall back to the host-native config cwd. let parent_turn_legacy_fallback_cwd = params - .parent_turn - .environments + .parent_environments .primary() .and_then(|environment| environment.cwd().to_abs_path().ok()) .unwrap_or_else(|| params.parent_turn.config.cwd.clone()); @@ -1157,6 +1159,8 @@ mod tests { async fn test_review_params() -> GuardianReviewSessionParams { let (session, turn) = crate::session::tests::make_session_and_context().await; + let session = Arc::new(session); + let parent_environments = session.services.turn_environments.snapshot().await; let model = turn.model_info.slug.clone(); let reasoning_effort = turn.reasoning_effort.clone(); let reasoning_summary = turn.reasoning_summary; @@ -1172,8 +1176,9 @@ mod tests { .expect("guardian config"); GuardianReviewSessionParams { - parent_session: Arc::new(session), + parent_session: session, parent_turn: Arc::new(turn), + parent_environments, spawn_config, request: GuardianApprovalRequest::Shell { id: "shell-1".to_string(), diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index b19cd9d35e7..9f448b5fc92 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -87,6 +87,12 @@ const GUARDIAN_MEMORY_CONTEXT_PROBE: &str = "guardian memory context probe"; const GUARDIAN_SKILL_NAME: &str = "guardian-context-probe"; const GUARDIAN_SKILL_BODY_PROBE: &str = "guardian skill body probe"; +async fn guardian_test_environments( + session: &Session, +) -> crate::environment_selection::TurnEnvironmentSnapshot { + session.services.turn_environments.snapshot().await +} + // The memories extension depends on codex-core, so this probe verifies the nested Guardian config // at request assembly without introducing a circular test dependency. struct GuardianMemoryContextEnabled(bool); @@ -1149,6 +1155,7 @@ async fn cancelled_guardian_review_emits_terminal_abort_without_warning() { let decision = review_approval_request_with_cancel( &session, &turn, + guardian_test_environments(&session).await, "review-cancelled-guardian".to_string(), GuardianApprovalRequest::ApplyPatch { id: "patch-1".to_string(), @@ -1469,6 +1476,7 @@ async fn guardian_request_model_for_auto_review( let (outcome, analytics_result) = run_guardian_review_session_for_test( Arc::clone(&session), turn, + guardian_test_environments(&session).await, GuardianApprovalRequest::Shell { id: "shell-1".to_string(), command: vec!["git".to_string(), "push".to_string()], @@ -1718,6 +1726,7 @@ async fn guardian_review_request_layout_matches_model_visible_request_snapshot() let outcome = run_guardian_review_session_for_test( Arc::clone(&session), Arc::clone(&turn), + guardian_test_environments(&session).await, request, Some("Sandbox denied outbound git push to github.com.".to_string()), guardian_output_schema(), @@ -1910,6 +1919,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: let first_outcome = run_guardian_review_session_for_test( Arc::clone(&session), Arc::clone(&turn), + guardian_test_environments(&session).await, first_request, Some("First retry reason".to_string()), guardian_output_schema(), @@ -1957,6 +1967,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: let second_outcome = run_guardian_review_session_for_test( Arc::clone(&session), Arc::clone(&turn), + guardian_test_environments(&session).await, second_request, Some("Second retry reason".to_string()), guardian_output_schema(), @@ -2000,6 +2011,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: let third_outcome = run_guardian_review_session_for_test( Arc::clone(&session), Arc::clone(&turn), + guardian_test_environments(&session).await, third_request, Some("Third retry reason".to_string()), guardian_output_schema(), @@ -2185,6 +2197,7 @@ async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow:: let first_outcome = run_guardian_review_session_for_test( Arc::clone(&session), Arc::clone(&turn), + guardian_test_environments(&session).await, GuardianApprovalRequest::Shell { id: "shell-1".to_string(), command: vec!["git".to_string(), "push".to_string()], @@ -2228,6 +2241,7 @@ async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow:: let second_outcome = run_guardian_review_session_for_test( Arc::clone(&session), Arc::clone(&turn), + guardian_test_environments(&session).await, GuardianApprovalRequest::Shell { id: "shell-2".to_string(), command: vec!["git".to_string(), "push".to_string()], @@ -2308,6 +2322,7 @@ async fn guardian_review_surfaces_responses_api_errors_in_rejection_reason() -> let decision = review_approval_request( &session, &turn, + guardian_test_environments(&session).await, "review-shell-guardian-error".to_string(), GuardianApprovalRequest::Shell { id: "shell-guardian-error".to_string(), @@ -2407,6 +2422,7 @@ async fn guardian_review_retries_transient_session_failure_then_approves() -> an let (outcome, metadata) = run_guardian_review_session_for_test( Arc::clone(&session), Arc::clone(&turn), + guardian_test_environments(&session).await, guardian_shell_request("shell-session-retry"), /*retry_reason*/ None, guardian_output_schema(), @@ -2448,6 +2464,7 @@ async fn guardian_review_does_not_retry_missing_assessment_payload() -> anyhow:: let decision = review_approval_request( &session, &turn, + guardian_test_environments(&session).await, "review-missing-assessment".to_string(), guardian_shell_request("shell-missing-assessment"), /*retry_reason*/ None, @@ -2498,6 +2515,7 @@ async fn guardian_review_retries_two_parse_failures_then_approves() -> anyhow::R let (outcome, metadata) = run_guardian_review_session_for_test( Arc::clone(&session), Arc::clone(&turn), + guardian_test_environments(&session).await, guardian_shell_request("shell-parse-retry"), /*retry_reason*/ None, guardian_output_schema(), @@ -2552,6 +2570,7 @@ async fn guardian_review_exhausts_three_failures_with_one_terminal_event() -> an let decision = review_approval_request( &session, &turn, + guardian_test_environments(&session).await, "review-exhausted-retry".to_string(), guardian_shell_request("shell-exhausted-retry"), /*retry_reason*/ None, @@ -2603,6 +2622,7 @@ async fn guardian_review_does_not_retry_valid_denial() -> anyhow::Result<()> { let decision = review_approval_request( &session, &turn, + guardian_test_environments(&session).await, "review-valid-denial".to_string(), guardian_shell_request("shell-valid-denial"), /*retry_reason*/ None, @@ -2706,6 +2726,7 @@ async fn guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history() -> review_approval_request( &session, &turn, + guardian_test_environments(&session).await, "review-shell-guardian-1".to_string(), initial_request, /*retry_reason*/ None @@ -2760,6 +2781,7 @@ async fn guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history() -> review_approval_request( &session_for_second, &turn_for_second, + guardian_test_environments(&session_for_second).await, "review-shell-guardian-2".to_string(), second_request, Some("trunk follow-up".to_string()), @@ -2807,6 +2829,7 @@ async fn guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history() -> let third_decision = review_approval_request( &session, &turn, + guardian_test_environments(&session).await, "review-shell-guardian-3".to_string(), third_request, Some("parallel follow-up".to_string()), diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 4dee0da717a..4be0507b819 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -12,6 +12,7 @@ //! inventory, so this module only handles the execution-time argument rewrite. use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use codex_api::OPENAI_FILE_UPLOAD_LIMIT_BYTES; use codex_api::upload_openai_file; @@ -22,6 +23,7 @@ use serde_json::Value as JsonValue; pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( sess: &Session, turn_context: &TurnContext, + step_context: &StepContext, arguments_value: Option, openai_file_input_params: Option<&[String]>, ) -> Result, String> { @@ -42,9 +44,14 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( let Some(value) = arguments.get(field_name) else { continue; }; - let Some(uploaded_value) = - rewrite_argument_value_for_openai_files(turn_context, auth.as_ref(), field_name, value) - .await? + let Some(uploaded_value) = rewrite_argument_value_for_openai_files( + turn_context, + step_context, + auth.as_ref(), + field_name, + value, + ) + .await? else { continue; }; @@ -60,6 +67,7 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( async fn rewrite_argument_value_for_openai_files( turn_context: &TurnContext, + step_context: &StepContext, auth: Option<&CodexAuth>, field_name: &str, value: &JsonValue, @@ -68,6 +76,7 @@ async fn rewrite_argument_value_for_openai_files( JsonValue::String(file_path) => { let rewritten = build_uploaded_argument_value( turn_context, + step_context, auth, field_name, /*index*/ None, @@ -84,6 +93,7 @@ async fn rewrite_argument_value_for_openai_files( }; let rewritten = build_uploaded_argument_value( turn_context, + step_context, auth, field_name, Some(index), @@ -100,6 +110,7 @@ async fn rewrite_argument_value_for_openai_files( async fn build_uploaded_argument_value( turn_context: &TurnContext, + step_context: &StepContext, auth: Option<&CodexAuth>, field_name: &str, index: Option, @@ -117,7 +128,7 @@ async fn build_uploaded_argument_value( if !auth.uses_codex_backend() { return Err("ChatGPT auth is required to upload files for Codex Apps tools".to_string()); } - let Some(turn_environment) = turn_context.environments.primary() else { + let Some(turn_environment) = step_context.environments.primary() else { return Err(contextualize_error( "no primary turn environment is available".to_string(), )); @@ -190,10 +201,9 @@ mod tests { use std::sync::Arc; use tempfile::tempdir; - fn set_primary_environment_cwd(turn_context: &mut TurnContext, cwd: &Path) { + fn set_primary_environment_cwd(step_context: &mut StepContext, cwd: &Path) { let cwd = AbsolutePathBuf::try_from(cwd).expect("absolute path"); - turn_context.permission_profile = codex_protocol::models::PermissionProfile::Disabled; - let primary = turn_context + let primary = step_context .environments .turn_environments .first_mut() @@ -209,6 +219,7 @@ mod tests { #[tokio::test] async fn openai_file_argument_rewrite_requires_declared_file_params() { let (session, turn_context) = make_session_and_context().await; + let step_context = StepContext::local_for_test(&turn_context); let arguments = Some(serde_json::json!({ "file": "/tmp/codex-smoke-file.txt" })); @@ -216,6 +227,7 @@ mod tests { let rewritten = rewrite_mcp_tool_arguments_for_openai_files( &session, &Arc::new(turn_context), + &step_context, arguments.clone(), /*openai_file_input_params*/ None, ) @@ -271,13 +283,15 @@ mod tests { .await; let (_, mut turn_context) = make_session_and_context().await; + let mut step_context = StepContext::local_for_test(&turn_context); + turn_context.permission_profile = codex_protocol::models::PermissionProfile::Disabled; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); let dir = tempdir().expect("temp dir"); let local_path = dir.path().join("file_report.csv"); tokio::fs::write(&local_path, b"hello") .await .expect("write local file"); - set_primary_environment_cwd(&mut turn_context, dir.path()); + set_primary_environment_cwd(&mut step_context, dir.path()); let mut config = (*turn_context.config).clone(); config.chatgpt_base_url = format!("{}/backend-api", server.uri()); @@ -285,6 +299,7 @@ mod tests { let rewritten = build_uploaded_argument_value( &turn_context, + &step_context, Some(&auth), "file", /*index*/ None, @@ -309,16 +324,19 @@ mod tests { #[tokio::test] async fn build_uploaded_argument_value_rejects_oversized_file_before_reading() { let (_, mut turn_context) = make_session_and_context().await; + let mut step_context = StepContext::local_for_test(&turn_context); + turn_context.permission_profile = codex_protocol::models::PermissionProfile::Disabled; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); let dir = tempdir().expect("temp dir"); let file_path = dir.path().join("oversized.bin"); let file = std::fs::File::create(&file_path).expect("create sparse file"); file.set_len(OPENAI_FILE_UPLOAD_LIMIT_BYTES + 1) .expect("size sparse file"); - set_primary_environment_cwd(&mut turn_context, dir.path()); + set_primary_environment_cwd(&mut step_context, dir.path()); let error = build_uploaded_argument_value( &turn_context, + &step_context, Some(&auth), "file", /*index*/ None, @@ -377,19 +395,22 @@ mod tests { .await; let (_, mut turn_context) = make_session_and_context().await; + let mut step_context = StepContext::local_for_test(&turn_context); + turn_context.permission_profile = codex_protocol::models::PermissionProfile::Disabled; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); let dir = tempdir().expect("temp dir"); let local_path = dir.path().join("file_report.csv"); tokio::fs::write(&local_path, b"hello") .await .expect("write local file"); - set_primary_environment_cwd(&mut turn_context, dir.path()); + set_primary_environment_cwd(&mut step_context, dir.path()); let mut config = (*turn_context.config).clone(); config.chatgpt_base_url = format!("{}/backend-api", server.uri()); turn_context.config = Arc::new(config); let rewritten = rewrite_argument_value_for_openai_files( &turn_context, + &step_context, Some(&auth), "file", &serde_json::json!("file_report.csv"), @@ -489,6 +510,8 @@ mod tests { .await; let (_, mut turn_context) = make_session_and_context().await; + let mut step_context = StepContext::local_for_test(&turn_context); + turn_context.permission_profile = codex_protocol::models::PermissionProfile::Disabled; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); let dir = tempdir().expect("temp dir"); tokio::fs::write(dir.path().join("one.csv"), b"one") @@ -497,13 +520,14 @@ mod tests { tokio::fs::write(dir.path().join("two.csv"), b"two") .await .expect("write second local file"); - set_primary_environment_cwd(&mut turn_context, dir.path()); + set_primary_environment_cwd(&mut step_context, dir.path()); let mut config = (*turn_context.config).clone(); config.chatgpt_base_url = format!("{}/backend-api", server.uri()); turn_context.config = Arc::new(config); let rewritten = rewrite_argument_value_for_openai_files( &turn_context, + &step_context, Some(&auth), "files", &serde_json::json!(["one.csv", "two.csv"]), @@ -537,12 +561,14 @@ mod tests { #[tokio::test] async fn rewrite_mcp_tool_arguments_for_openai_files_surfaces_upload_failures() { let (mut session, turn_context) = make_session_and_context().await; + let step_context = StepContext::local_for_test(&turn_context); session.services.auth_manager = crate::test_support::auth_manager_from_auth( CodexAuth::create_dummy_chatgpt_auth_for_testing(), ); let error = rewrite_mcp_tool_arguments_for_openai_files( &session, &turn_context, + &step_context, Some(serde_json::json!({ "file": "/definitely/missing/file.csv", })), diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index e8cbb56c551..73f33d5b143 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -19,6 +19,7 @@ use crate::mcp_openai_file::rewrite_mcp_tool_arguments_for_openai_files; use crate::mcp_tool_approval_templates::RenderedMcpToolApprovalParam; use crate::mcp_tool_approval_templates::render_mcp_tool_approval_template; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::hook_names::HookToolName; use crate::tools::sandboxing::PermissionRequestPayload; @@ -108,9 +109,11 @@ const MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES: usize = DEFAULT_OUTPUT_BYTES_CAP; /// Handles the specified tool call and dispatches the appropriate MCP tool-call /// item lifecycle events to the `Session`. +#[allow(clippy::too_many_arguments)] pub(crate) async fn handle_mcp_tool_call( sess: Arc, turn_context: &Arc, + step_context: &Arc, call_id: String, server: String, tool_name: String, @@ -228,6 +231,7 @@ pub(crate) async fn handle_mcp_tool_call( if let Some(decision) = maybe_request_mcp_tool_approval( &sess, turn_context, + step_context, &call_id, &invocation, &hook_tool_name, @@ -243,6 +247,7 @@ pub(crate) async fn handle_mcp_tool_call( return handle_approved_mcp_tool_call( sess.as_ref(), turn_context.as_ref(), + step_context.as_ref(), &call_id, invocation, metadata.as_ref(), @@ -298,6 +303,7 @@ pub(crate) async fn handle_mcp_tool_call( handle_approved_mcp_tool_call( sess.as_ref(), turn_context.as_ref(), + step_context.as_ref(), &call_id, invocation, metadata.as_ref(), @@ -320,6 +326,7 @@ struct McpToolCallItemMetadata { async fn handle_approved_mcp_tool_call( sess: &Session, turn_context: &TurnContext, + step_context: &StepContext, call_id: &str, invocation: McpInvocation, metadata: Option<&McpToolApprovalMetadata>, @@ -342,6 +349,7 @@ async fn handle_approved_mcp_tool_call( let rewrite = rewrite_mcp_tool_arguments_for_openai_files( sess, turn_context, + step_context, arguments_value.clone(), metadata.and_then(|metadata| metadata.openai_file_input_params.as_deref()), ) @@ -1160,9 +1168,11 @@ fn mcp_tool_approval_prompt_options( } } +#[allow(clippy::too_many_arguments)] async fn maybe_request_mcp_tool_approval( sess: &Arc, turn_context: &Arc, + step_context: &Arc, call_id: &str, invocation: &McpInvocation, hook_tool_name: &HookToolName, @@ -1238,6 +1248,7 @@ async fn maybe_request_mcp_tool_approval( let decision = review_approval_request( sess, turn_context, + step_context.environments.clone(), review_id.clone(), build_guardian_mcp_tool_review_request(call_id, invocation, metadata), /*retry_reason*/ None, @@ -1344,6 +1355,32 @@ async fn maybe_request_mcp_tool_approval( Some(decision) } +#[cfg(test)] +async fn maybe_request_mcp_tool_approval_for_test( + sess: &Arc, + turn_context: &Arc, + call_id: &str, + invocation: &McpInvocation, + hook_tool_name: &HookToolName, + metadata: Option<&McpToolApprovalMetadata>, + approval_mode: AppToolApproval, +) -> Option { + let step_context = Arc::new(StepContext::new( + sess.services.turn_environments.snapshot().await, + )); + maybe_request_mcp_tool_approval( + sess, + turn_context, + &step_context, + call_id, + invocation, + hook_tool_name, + metadata, + approval_mode, + ) + .await +} + pub(crate) fn mcp_approvals_reviewer( turn_context: &TurnContext, server_name: &str, diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 56d1ff9b6cd..30d60b1c051 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -2300,7 +2300,7 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() { openai_file_input_params: None, }; - let decision = maybe_request_mcp_tool_approval( + let decision = maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-1", @@ -2374,7 +2374,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { openai_file_input_params: None, }; - let decision = maybe_request_mcp_tool_approval( + let decision = maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-guardian", @@ -2431,7 +2431,7 @@ async fn permission_request_hook_allows_mcp_tool_call() { openai_file_input_params: None, }; - let decision = maybe_request_mcp_tool_approval( + let decision = maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-mcp-hook", @@ -2493,7 +2493,7 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() { arguments: Some(serde_json::json!({ "entities": [] })), }; - let decision = maybe_request_mcp_tool_approval( + let decision = maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-mcp-hook-no-metadata", @@ -2573,7 +2573,7 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() { let session = Arc::new(session); let turn_context = Arc::new(turn_context); - let decision = maybe_request_mcp_tool_approval( + let decision = maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-mcp-remembered", @@ -2654,7 +2654,7 @@ async fn guardian_mode_mcp_denial_returns_rationale_message() { openai_file_input_params: None, }; - let decision = maybe_request_mcp_tool_approval( + let decision = maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-guardian-deny", @@ -2712,7 +2712,7 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); tokio::spawn(async move { - maybe_request_mcp_tool_approval( + maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-prompt", @@ -2768,7 +2768,7 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() { AppToolApproval::Prompt, AppToolApproval::Approve, ] { - let decision = maybe_request_mcp_tool_approval( + let decision = maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-2", @@ -2856,7 +2856,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { let session = Arc::new(session); let turn_context = Arc::new(turn_context); - let decision = maybe_request_mcp_tool_approval( + let decision = maybe_request_mcp_tool_approval_for_test( &session, &turn_context, "call-3", diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 0da42d8a4e3..f6a87c6e597 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -76,8 +76,12 @@ pub(crate) async fn build_prompt_input_from_session( input: Vec, ) -> CodexResult> { let turn_context = sess.new_default_turn().await; - sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref()) - .await; + let step_context = sess.prepare_step_for_request().await; + sess.record_context_updates_and_set_reference_context_item( + turn_context.as_ref(), + step_context.as_ref(), + ) + .await; if !input.is_empty() { let response_item = sess.response_item_from_user_input(turn_context.as_ref(), input); @@ -89,7 +93,13 @@ pub(crate) async fn build_prompt_input_from_session( .clone_history() .await .for_prompt(&turn_context.model_info.input_modalities); - let router = built_tools(sess, turn_context.as_ref(), &CancellationToken::new()).await?; + let router = built_tools( + sess, + turn_context.as_ref(), + step_context, + &CancellationToken::new(), + ) + .await?; let base_instructions = sess.get_base_instructions().await; let prompt = build_prompt( prompt_input, diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 4ed79d04e38..6b15c8dbde3 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -237,8 +237,9 @@ pub(super) async fn user_input_or_turn_inner( .set_responsesapi_client_metadata(responsesapi_client_metadata); } current_context.session_telemetry.user_prompt(&items); + // Refresh outside the task so cancellation cannot consume it midway. sess.refresh_mcp_servers_if_requested( - ¤t_context, + current_context.as_ref(), Some(sess.mcp_elicitation_reviewer()), ) .await; @@ -258,7 +259,7 @@ pub(super) async fn user_input_or_turn_inner( }); } sess.spawn_task( - Arc::clone(¤t_context), + current_context, task_input, crate::tasks::RegularTask::new(), ) @@ -310,12 +311,8 @@ pub async fn run_user_shell_command(sess: &Arc, sub_id: String, command } let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; - sess.spawn_task( - Arc::clone(&turn_context), - Vec::new(), - UserShellCommandTask::new(command), - ) - .await; + sess.spawn_task(turn_context, Vec::new(), UserShellCommandTask::new(command)) + .await; } pub async fn resolve_elicitation( @@ -444,8 +441,7 @@ pub async fn reload_user_config(sess: &Arc) { pub async fn compact(sess: &Arc, sub_id: String) { let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; - sess.spawn_task(Arc::clone(&turn_context), Vec::new(), CompactTask) - .await; + sess.spawn_task(turn_context, Vec::new(), CompactTask).await; } pub async fn thread_rollback(sess: &Arc, sub_id: String, num_turns: u32) { @@ -670,13 +666,15 @@ pub async fn review( .await; sess.refresh_mcp_servers_if_requested(&turn_context, Some(sess.mcp_elicitation_reviewer())) .await; - #[allow(deprecated)] - match resolve_review_request(review_request, &turn_context.cwd) { + let step_context = sess.prepare_step_for_request().await; + let review_cwd = step_context.effective_cwd(turn_context.as_ref()); + match resolve_review_request(review_request, &review_cwd) { Ok(resolved) => { spawn_review_thread( Arc::clone(sess), Arc::clone(config), - turn_context.clone(), + turn_context, + step_context.environments.clone(), sub_id, resolved, ) diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index d1b89822232..407d13c2617 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -320,10 +320,10 @@ impl Session { let environment_manager = self.services.turn_environments.environment_manager(); // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd // values can be used without falling back to the legacy host cwd. - let cwd = turn_context - .environments + let environments = self.services.turn_environments.snapshot().await; + let cwd = environments .primary() - .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) + .and_then(|environment| environment.cwd().to_abs_path().ok()) .map(|cwd| cwd.to_path_buf()) .unwrap_or_else(|| { #[allow(deprecated)] @@ -492,9 +492,14 @@ async fn review_guardian_mcp_elicitation( }; let review_id = crate::guardian::new_guardian_review_id(); + // V0 only allows the selected environment to move from starting to attached, so use + // its latest snapshot. Bind elicitations to their originating MCP call before + // supporting selection changes or detachment during a turn. + let environments = session.services.turn_environments.snapshot().await; let decision = crate::guardian::review_approval_request( &session, &turn_context, + environments, review_id.clone(), guardian_request, /*retry_reason*/ None, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index cf191235919..dbbdf25b58a 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -212,6 +212,7 @@ mod review; mod rollout_reconstruction; #[allow(clippy::module_inception)] pub(crate) mod session; +pub(crate) mod step_context; mod token_budget; pub(crate) mod turn; pub(crate) mod turn_context; @@ -228,6 +229,7 @@ use self::session::AppServerClientMetadata; use self::session::Session; use self::session::SessionConfiguration; pub(crate) use self::session::SessionSettingsUpdate; +use self::step_context::StepContext; #[cfg(test)] use self::turn::AssistantMessageStreamParsers; #[cfg(test)] @@ -374,7 +376,6 @@ use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnModerationMetadataEvent; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; -use codex_tools::ToolEnvironmentMode; use codex_tools::UnifiedExecShellMode; use codex_utils_absolute_path::AbsolutePathBuf; #[cfg(test)] @@ -1623,6 +1624,7 @@ impl Session { &self, reference_context_item: Option<&TurnContextItem>, current_context: &TurnContext, + step_context: &StepContext, ) -> Vec { // TODO: Make context updates a pure diff of persisted previous/current TurnContextItem // state so replay/backtracking is deterministic. Runtime inputs that affect model-visible @@ -1638,12 +1640,24 @@ impl Session { reference_context_item, previous_turn_settings.as_ref(), current_context, + step_context, shell.as_ref(), exec_policy.as_ref(), self.features.enabled(Feature::Personality), ) } + #[cfg(test)] + async fn build_settings_update_items_for_test( + &self, + reference_context_item: Option<&TurnContextItem>, + current_context: &TurnContext, + ) -> Vec { + let step_context = StepContext::new(self.services.turn_environments.snapshot().await); + self.build_settings_update_items(reference_context_item, current_context, &step_context) + .await + } + /// Record a terminal CodexErr before the app-server completion notification is reduced. pub(crate) fn track_turn_codex_error(&self, turn_context: &TurnContext, error: &CodexErr) { self.services @@ -2156,6 +2170,7 @@ impl Session { pub(crate) async fn request_permissions_for_environment( self: &Arc, turn_context: &Arc, + environments: &TurnEnvironmentSnapshot, call_id: String, args: RequestPermissionsArgs, environment: TurnEnvironmentSelection, @@ -2206,6 +2221,7 @@ impl Session { let review_id = crate::guardian::new_guardian_review_id(); let session = Arc::clone(self); let turn = Arc::clone(turn_context); + let environments = environments.clone(); let request = crate::guardian::GuardianApprovalRequest::RequestPermissions { id: call_id, turn_id: turn_context.sub_id.clone(), @@ -2215,6 +2231,7 @@ impl Session { let review_rx = crate::guardian::spawn_approval_request_review( session, turn, + environments, review_id, request, /*retry_reason*/ None, @@ -2324,18 +2341,18 @@ impl Session { pub(crate) async fn request_permissions_for_cwd( self: &Arc, turn_context: &Arc, + environments: &TurnEnvironmentSnapshot, call_id: String, args: RequestPermissionsArgs, cwd: AbsolutePathBuf, cancellation_token: CancellationToken, ) -> Option { let turn_environment = match args.environment_id.as_deref() { - Some(environment_id) => turn_context - .environments + Some(environment_id) => environments .turn_environments .iter() .find(|environment| environment.environment_id == environment_id), - None => turn_context.environments.primary(), + None => environments.primary(), }; let Some(turn_environment) = turn_environment else { return Some(RequestPermissionsResponse { @@ -2348,6 +2365,7 @@ impl Session { environment.cwd = PathUri::from_abs_path(&cwd); self.request_permissions_for_environment( turn_context, + environments, call_id, args, environment, @@ -2863,6 +2881,7 @@ impl Session { pub(crate) async fn build_initial_context( &self, turn_context: &TurnContext, + step_context: &StepContext, ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::with_capacity(2); @@ -2894,14 +2913,14 @@ impl Session { developer_sections.push(model_switch_message); } if turn_context.config.include_permissions_instructions { + let cwd = step_context.effective_cwd(turn_context); developer_sections.push( PermissionsInstructions::from_permission_profile( &turn_context.permission_profile, turn_context.approval_policy.value(), turn_context.config.approvals_reviewer, self.services.exec_policy.current().as_ref(), - #[allow(deprecated)] - &turn_context.cwd, + &cwd, turn_context .config .features @@ -3071,9 +3090,13 @@ impl Session { .format_environment_context_subagents(self.thread_id) .await; contextual_user_sections.push( - crate::context::EnvironmentContext::from_turn_context(turn_context, shell.as_ref()) - .with_subagents(subagents) - .render(), + crate::context::EnvironmentContext::from_turn_context( + turn_context, + step_context, + shell.as_ref(), + ) + .with_subagents(subagents) + .render(), ); } @@ -3121,6 +3144,16 @@ impl Session { items } + #[cfg(test)] + pub(crate) async fn build_initial_context_for_test( + &self, + turn_context: &TurnContext, + ) -> Vec { + let step_context = StepContext::new(self.services.turn_environments.snapshot().await); + self.build_initial_context(turn_context, &step_context) + .await + } + pub(crate) async fn persist_rollout_items(&self, items: &[RolloutItem]) { if let Some(live_thread) = self.live_thread() && let Err(e) = live_thread.append_items(items).await @@ -3154,13 +3187,14 @@ impl Session { pub(crate) async fn maybe_start_new_context_window( &self, turn_context: &TurnContext, + step_context: &StepContext, ) -> Option { let window_id = { let mut state = self.state.lock().await; state.start_new_context_window_if_requested() }; let window_id = window_id?; - let context_items = self.build_initial_context(turn_context).await; + let context_items = self.build_initial_context(turn_context, step_context).await; let turn_context_item = turn_context.to_turn_context_item(); let replacement_history = context_items; { @@ -3206,6 +3240,7 @@ impl Session { pub(crate) async fn record_context_updates_and_set_reference_context_item( &self, turn_context: &TurnContext, + step_context: &StepContext, ) { let reference_context_item = { let state = self.state.lock().await; @@ -3213,11 +3248,15 @@ impl Session { }; let should_inject_full_context = reference_context_item.is_none(); let context_items = if should_inject_full_context { - self.build_initial_context(turn_context).await + self.build_initial_context(turn_context, step_context).await } else { // Steady-state path: append only context diffs to minimize token overhead. - self.build_settings_update_items(reference_context_item.as_ref(), turn_context) - .await + self.build_settings_update_items( + reference_context_item.as_ref(), + turn_context, + step_context, + ) + .await }; let turn_context_item = turn_context.to_turn_context_item(); if !context_items.is_empty() { @@ -3235,6 +3274,16 @@ impl Session { state.set_reference_context_item(Some(turn_context_item)); } + #[cfg(test)] + pub(crate) async fn record_context_updates_and_set_reference_context_item_for_test( + &self, + turn_context: &TurnContext, + ) { + let step_context = StepContext::new(self.services.turn_environments.snapshot().await); + self.record_context_updates_and_set_reference_context_item(turn_context, &step_context) + .await; + } + pub(crate) async fn update_token_usage_info( &self, turn_context: &TurnContext, diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index bbdfc3f984c..a58ef48a914 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -6,6 +6,7 @@ pub(super) async fn spawn_review_thread( sess: Arc, config: Arc, parent_turn_context: Arc, + parent_environments: TurnEnvironmentSnapshot, sub_id: String, resolved: crate::review_prompts::ResolvedReviewRequest, ) { @@ -114,7 +115,6 @@ pub(super) async fn spawn_review_thread( reasoning_summary, session_source, parent_thread_id: parent_turn_context.parent_thread_id, - environments: parent_turn_context.environments.clone(), available_models, unified_exec_shell_mode, current_date: parent_turn_context.current_date.clone(), @@ -151,20 +151,30 @@ pub(super) async fn spawn_review_thread( }], client_id: None, }]; - let tc = Arc::new(review_turn_context); - if tc.environments.single_local_environment_cwd().is_some() { - tc.turn_metadata_state.spawn_git_enrichment_task(); + let review_context = Arc::new(review_turn_context); + if parent_environments.single_local_environment_cwd().is_some() { + review_context + .turn_metadata_state + .spawn_git_enrichment_task(); } // TODO(ccunningham): Review turns currently rely on `spawn_task` for TurnComplete but do not // emit a parent TurnStarted. Consider giving review a full parent turn lifecycle // (TurnStarted + TurnComplete) for consistency with other standalone tasks. - sess.spawn_task(tc.clone(), input, ReviewTask::new()).await; + sess.spawn_task( + Arc::clone(&review_context), + input, + ReviewTask::new(parent_environments), + ) + .await; // Announce entering review mode so UIs can switch modes. let review_request = ReviewRequest { target: resolved.target, user_facing_hint: Some(resolved.user_facing_hint), }; - sess.send_event(&tc, EventMsg::EnteredReviewMode(review_request)) - .await; + sess.send_event( + review_context.as_ref(), + EventMsg::EnteredReviewMode(review_request), + ) + .await; } diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs new file mode 100644 index 00000000000..923315c4319 --- /dev/null +++ b/codex-rs/core/src/session/step_context.rs @@ -0,0 +1,67 @@ +use std::sync::Arc; + +use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::session::session::Session; +use crate::session::turn_context::TurnContext; +use codex_tools::ToolEnvironmentMode; +use codex_utils_absolute_path::AbsolutePathBuf; + +/// Immutable environment state shared by model context, tools, and tool calls. +/// Persisting and diffing starting/attached state remain deferred to a dedicated step baseline. +#[derive(Debug)] +pub(crate) struct StepContext { + pub(crate) environments: TurnEnvironmentSnapshot, +} + +impl StepContext { + pub(crate) fn new(environments: TurnEnvironmentSnapshot) -> Self { + Self { environments } + } + + pub(crate) fn tool_environment_mode(&self) -> ToolEnvironmentMode { + ToolEnvironmentMode::from_count(self.environments.turn_environments.len()) + } + + pub(crate) fn effective_cwd(&self, turn: &TurnContext) -> AbsolutePathBuf { + self.environments + .primary() + .map(super::turn_context::TurnEnvironment::cwd) + .or_else(|| { + self.environments + .starting + .first() + .map(|environment| &environment.selection.cwd) + }) + .and_then(|cwd| cwd.to_abs_path().ok()) + .unwrap_or_else(|| { + #[allow(deprecated)] + turn.cwd.clone() + }) + } + + #[cfg(test)] + pub(crate) fn local_for_test(turn: &TurnContext) -> Self { + #[allow(deprecated)] + let cwd = codex_utils_path_uri::PathUri::from_abs_path(&turn.cwd); + let environment = Arc::new( + codex_exec_server::Environment::create_for_tests(/*exec_server_url*/ None) + .expect("create local test environment"), + ); + Self::new(TurnEnvironmentSnapshot::from_turn_environments(vec![ + crate::session::turn_context::TurnEnvironment::new( + codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + environment, + cwd, + /*shell*/ None, + ), + ])) + } +} + +impl Session { + pub(crate) async fn prepare_step_for_request(&self) -> Arc { + Arc::new(StepContext::new( + self.services.turn_environments.snapshot().await, + )) + } +} diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 84c4a38887b..544fd22deb4 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -568,9 +568,18 @@ async fn preview_session_start_hooks( ) } -fn test_tool_runtime(session: Arc, turn_context: Arc) -> ToolCallRuntime { - let router = Arc::new(ToolRouter::from_turn_context( +pub(crate) async fn step_context_for_session(session: &Session) -> Arc { + session.prepare_step_for_request().await +} + +async fn test_tool_runtime( + session: Arc, + turn_context: Arc, +) -> ToolCallRuntime { + let step_context = step_context_for_session(session.as_ref()).await; + let router = Arc::new(ToolRouter::from_contexts( &turn_context, + step_context, crate::tools::router::ToolRouterParams { tool_suggest_candidates: None, mcp_tools: None, @@ -1049,12 +1058,14 @@ async fn danger_full_access_tool_attempts_do_not_enforce_managed_network() -> an let turn = session.new_default_turn().await; assert!(turn.network.is_none()); + let step = step_context_for_session(session.as_ref()).await; let mut orchestrator = crate::tools::orchestrator::ToolOrchestrator::new(); let mut tool = ProbeToolRuntime::default(); let tool_ctx = crate::tools::sandboxing::ToolCtx { session: Arc::clone(&session), turn: Arc::clone(&turn), + environments: step.environments.clone(), call_id: "probe-call".to_string(), tool_name: codex_tools::ToolName::plain("probe"), }; @@ -1875,15 +1886,15 @@ async fn resumed_history_injects_initial_context_on_first_context_update_only() assert_eq!(expected, history_before_seed.raw_items()); session - .record_context_updates_and_set_reference_context_item(&turn_context) + .record_context_updates_and_set_reference_context_item_for_test(&turn_context) .await; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; expected.extend(initial_context); let history_after_seed = session.clone_history().await; assert_eq!(expected, history_after_seed.raw_items()); session - .record_context_updates_and_set_reference_context_item(&turn_context) + .record_context_updates_and_set_reference_context_item_for_test(&turn_context) .await; let history_after_second_seed = session.clone_history().await; assert_eq!( @@ -2759,7 +2770,7 @@ async fn thread_rollback_drops_last_turn_from_history() { ) .await; - let initial_context = sess.build_initial_context(tc.as_ref()).await; + let initial_context = sess.build_initial_context_for_test(tc.as_ref()).await; let turn_1 = vec![ user_message("turn 1 user"), assistant_message("turn 1 assistant"), @@ -2827,7 +2838,7 @@ async fn thread_rollback_clears_history_when_num_turns_exceeds_existing_turns() ) .await; - let initial_context = sess.build_initial_context(tc.as_ref()).await; + let initial_context = sess.build_initial_context_for_test(tc.as_ref()).await; let turn_1 = vec![user_message("turn 1 user")]; let mut full_history = Vec::new(); full_history.extend(initial_context.clone()); @@ -2853,7 +2864,7 @@ async fn thread_rollback_clears_history_when_num_turns_exceeds_existing_turns() async fn thread_rollback_fails_without_persisted_thread_history() { let (sess, tc, rx) = make_session_and_context_with_rx().await; - let initial_context = sess.build_initial_context(tc.as_ref()).await; + let initial_context = sess.build_initial_context_for_test(tc.as_ref()).await; sess.record_conversation_items(tc.as_ref(), &initial_context) .await; @@ -3242,7 +3253,7 @@ async fn thread_rollback_persists_marker_and_replays_cumulatively() { async fn thread_rollback_fails_when_turn_in_progress() { let (sess, tc, rx) = make_session_and_context_with_rx().await; - let initial_context = sess.build_initial_context(tc.as_ref()).await; + let initial_context = sess.build_initial_context_for_test(tc.as_ref()).await; sess.record_conversation_items(tc.as_ref(), &initial_context) .await; @@ -3263,7 +3274,7 @@ async fn thread_rollback_fails_when_turn_in_progress() { async fn thread_rollback_fails_when_num_turns_is_zero() { let (sess, tc, rx) = make_session_and_context_with_rx().await; - let initial_context = sess.build_initial_context(tc.as_ref()).await; + let initial_context = sess.build_initial_context_for_test(tc.as_ref()).await; sess.record_conversation_items(tc.as_ref(), &initial_context) .await; @@ -4764,7 +4775,16 @@ async fn absolute_cwd_update_with_turn_environment_is_allowed() { let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd, absolute_cwd); assert_eq!(turn_context.config.cwd, absolute_cwd); - assert_eq!(turn_context.environments.turn_environments.len(), 1); + assert_eq!( + session + .services + .turn_environments + .snapshot() + .await + .turn_environments + .len(), + 1 + ); } #[tokio::test] @@ -5071,7 +5091,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { model_info, &models_manager, /*network*/ None, - resolved_turn_environments, session_configuration.cwd().clone(), "turn_id".to_string(), skills_snapshot, @@ -5601,8 +5620,7 @@ async fn request_permissions_emits_event_when_granular_policy_allows_requests() })) .expect("test setup should allow updating approval policy"); - let session = Arc::new(session); - let turn_context = Arc::new(turn_context); + let step_context = step_context_for_session(session.as_ref()).await; let call_id = "call-1".to_string(); let expected_response = codex_protocol::request_permissions::RequestPermissionsResponse { permissions: RequestPermissionProfile { @@ -5618,9 +5636,10 @@ async fn request_permissions_emits_event_when_granular_policy_allows_requests() let handle = tokio::spawn({ let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); + let step_context = Arc::clone(&step_context); let call_id = call_id.clone(); async move { - let environment = turn_context + let environment = step_context .environments .primary() .expect("primary environment") @@ -5628,6 +5647,7 @@ async fn request_permissions_emits_event_when_granular_policy_allows_requests() session .request_permissions_for_environment( &turn_context, + &step_context.environments, call_id, codex_protocol::request_permissions::RequestPermissionsArgs { environment_id: None, @@ -5677,6 +5697,7 @@ async fn request_permissions_emits_event_when_granular_policy_allows_requests() #[tokio::test] async fn request_permissions_tool_resolves_relative_paths_against_selected_environment() { let (session, mut turn_context, rx) = make_session_and_context_with_rx().await; + let mut step_context = step_context_for_session(session.as_ref()).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let environment_cwd = { #[allow(deprecated)] @@ -5695,8 +5716,9 @@ async fn request_permissions_tool_resolves_relative_paths_against_selected_envir mcp_elicitations: true, })) .expect("test setup should allow updating approval policy"); - let current_environment = turn_context_mut.environments.turn_environments[0].clone(); - turn_context_mut.environments.turn_environments[0] = TurnEnvironment::new( + let step_context_mut = Arc::get_mut(&mut step_context).expect("single step ref"); + let current_environment = step_context_mut.environments.turn_environments[0].clone(); + step_context_mut.environments.turn_environments[0] = TurnEnvironment::new( "remote".to_string(), current_environment.environment, PathUri::from_abs_path(&environment_cwd), @@ -5709,6 +5731,7 @@ async fn request_permissions_tool_resolves_relative_paths_against_selected_envir let handle = tokio::spawn({ let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); + let step_context = Arc::clone(&step_context); let tracker = Arc::clone(&tracker); let call_id = call_id.clone(); async move { @@ -5716,6 +5739,7 @@ async fn request_permissions_tool_resolves_relative_paths_against_selected_envir .handle(ToolInvocation { session, turn: turn_context, + step: step_context, cancellation_token: CancellationToken::new(), tracker, call_id, @@ -5786,10 +5810,13 @@ async fn request_permissions_tool_resolves_relative_paths_against_selected_envir #[tokio::test] async fn request_permissions_tool_rejects_unknown_environment_id() { let (session, turn_context) = make_session_and_context().await; + let session = Arc::new(session); + let step_context = step_context_for_session(session.as_ref()).await; let result = RequestPermissionsHandler .handle(ToolInvocation { - session: Arc::new(session), + session, turn: Arc::new(turn_context), + step: step_context, cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "call-1".to_string(), @@ -5831,8 +5858,7 @@ async fn request_permissions_response_materializes_session_cwd_grants_before_rec })) .expect("test setup should allow updating approval policy"); - let session = Arc::new(session); - let turn_context = Arc::new(turn_context); + let step_context = step_context_for_session(session.as_ref()).await; let call_id = "call-1".to_string(); let requested_permissions = RequestPermissionProfile { file_system: Some(FileSystemPermissions { @@ -5850,10 +5876,11 @@ async fn request_permissions_response_materializes_session_cwd_grants_before_rec let handle = tokio::spawn({ let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); + let step_context = Arc::clone(&step_context); let call_id = call_id.clone(); let requested_permissions = requested_permissions.clone(); async move { - let environment = turn_context + let environment = step_context .environments .primary() .expect("primary environment") @@ -5861,6 +5888,7 @@ async fn request_permissions_response_materializes_session_cwd_grants_before_rec session .request_permissions_for_environment( &turn_context, + &step_context.environments, call_id, codex_protocol::request_permissions::RequestPermissionsArgs { environment_id: None, @@ -5941,10 +5969,9 @@ async fn request_permissions_is_auto_denied_when_granular_policy_blocks_tool_req })) .expect("test setup should allow updating approval policy"); - let session = Arc::new(session); - let turn_context = Arc::new(turn_context); + let step_context = step_context_for_session(session.as_ref()).await; let call_id = "call-1".to_string(); - let environment = turn_context + let environment = step_context .environments .primary() .expect("primary environment") @@ -5952,6 +5979,7 @@ async fn request_permissions_is_auto_denied_when_granular_policy_blocks_tool_req let response = session .request_permissions_for_environment( &turn_context, + &step_context.environments, call_id, codex_protocol::request_permissions::RequestPermissionsArgs { environment_id: None, @@ -6216,17 +6244,16 @@ async fn turn_environments_set_primary_environment() { .await .expect("turn should start"); - let turn_environments = &turn_context.environments; + let turn_environments = session.services.turn_environments.snapshot().await; assert_eq!(turn_environments.turn_environments.len(), 1); - let turn_environment = turn_context - .environments + let turn_environment = turn_environments .primary() .expect("primary environment should be set"); assert!(std::sync::Arc::ptr_eq( &turn_environment.environment, &turn_environments.turn_environments[0].environment )); - assert!(!turn_context.environments.turn_environments.is_empty()); + assert!(!turn_environments.turn_environments.is_empty()); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd.as_path(), selected_cwd.as_path()); @@ -6247,14 +6274,15 @@ async fn turn_environments_set_primary_environment() { )); let default_turn = session.new_default_turn().await; + let default_environments = session.services.turn_environments.snapshot().await; assert!(Arc::ptr_eq( &stored_environment, - &default_turn - .environments + &default_environments .primary() .expect("default turn primary environment") .environment )); + assert_eq!(default_turn.config.cwd.as_path(), selected_cwd.as_path()); } #[tokio::test] @@ -6274,10 +6302,9 @@ async fn default_turn_does_not_overlay_legacy_fallback_cwd_onto_stored_thread_en let turn_context = session.new_default_turn().await; - let turn_environments = &turn_context.environments; + let turn_environments = session.services.turn_environments.snapshot().await; assert_eq!(turn_environments.turn_environments.len(), 1); - let turn_environment = turn_context - .environments + let turn_environment = turn_environments .primary() .expect("primary environment should be set"); assert!(std::sync::Arc::ptr_eq( @@ -6303,43 +6330,40 @@ async fn default_turn_honors_empty_stored_thread_environments() { let turn_context = session.new_default_turn().await; - assert!(turn_context.environments.primary().is_none()); - assert!(turn_context.environments.turn_environments.is_empty()); + let environments = session.services.turn_environments.snapshot().await; + assert!(environments.primary().is_none()); + assert!(environments.turn_environments.is_empty()); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd, session_cwd); assert_eq!(turn_context.config.cwd, session_cwd); - assert_eq!(turn_context.environments.turn_environments.len(), 0); + assert_eq!(environments.turn_environments.len(), 0); } #[tokio::test] async fn primary_environment_uses_first_turn_environment() { - let (_session, mut turn_context) = make_session_and_context().await; - let first_environment = turn_context.environments.turn_environments[0].clone(); + let (session, turn_context) = make_session_and_context().await; + let mut environments = session.services.turn_environments.snapshot().await; + let first_environment = environments.turn_environments[0].clone(); #[allow(deprecated)] let second_cwd = turn_context.cwd.join("second"); let second_cwd_uri = codex_utils_path_uri::PathUri::from_abs_path(&second_cwd); - turn_context - .environments - .turn_environments - .push(TurnEnvironment::new( - "second".to_string(), - Arc::clone(&first_environment.environment), - second_cwd_uri.clone(), - /*shell*/ None, - )); + environments.turn_environments.push(TurnEnvironment::new( + "second".to_string(), + Arc::clone(&first_environment.environment), + second_cwd_uri.clone(), + /*shell*/ None, + )); assert_eq!( - turn_context - .environments + environments .primary() .expect("primary environment") .environment_id, first_environment.environment_id ); assert_eq!( - turn_context - .environments + environments .turn_environments .iter() .find(|environment| environment.environment_id == "second") @@ -6347,11 +6371,8 @@ async fn primary_environment_uses_first_turn_environment() { .cwd(), &second_cwd_uri ); - assert_eq!(turn_context.environments.turn_environments.len(), 2); - assert_eq!( - turn_context.environments.turn_environments[1].cwd(), - &second_cwd_uri - ); + assert_eq!(environments.turn_environments.len(), 2); + assert_eq!(environments.turn_environments[1].cwd(), &second_cwd_uri); } #[tokio::test] @@ -6372,8 +6393,9 @@ async fn empty_turn_environments_clear_primary_environment() { .await .expect("turn should start"); - assert!(turn_context.environments.primary().is_none()); - assert!(turn_context.environments.turn_environments.is_empty()); + let environments = session.services.turn_environments.snapshot().await; + assert!(environments.primary().is_none()); + assert!(environments.turn_environments.is_empty()); #[allow(deprecated)] let turn_cwd = turn_context.cwd.clone(); assert_eq!(turn_cwd, session.get_config().await.cwd); @@ -7114,7 +7136,6 @@ where model_info, &models_manager, /*network*/ None, - resolved_turn_environments, session_configuration.cwd().clone(), "turn_id".to_string(), skills_snapshot, @@ -7289,7 +7310,7 @@ async fn build_settings_update_items_emits_environment_item_for_network_changes( let reference_context_item = previous_context.to_turn_context_item(); let update_items = session - .build_settings_update_items(Some(&reference_context_item), ¤t_context) + .build_settings_update_items_for_test(Some(&reference_context_item), ¤t_context) .await; let environment_update = user_input_texts(&update_items) @@ -7303,18 +7324,20 @@ async fn build_settings_update_items_emits_environment_item_for_network_changes( #[tokio::test] async fn environment_context_uses_session_shell_when_environment_shell_is_absent() { - let (mut session, mut turn_context) = make_session_and_context().await; + let (mut session, turn_context) = make_session_and_context().await; + let mut step_context = StepContext::new(session.services.turn_environments.snapshot().await); session.services.user_shell = Arc::new(crate::shell::Shell { shell_type: crate::shell::ShellType::PowerShell, shell_path: PathBuf::from("powershell"), }); - for environment in &mut turn_context.environments.turn_environments { + for environment in &mut step_context.environments.turn_environments { environment.shell = None; } let session_shell = session.user_shell(); let environment_context = crate::context::EnvironmentContext::from_turn_context( &turn_context, + &step_context, session_shell.as_ref(), ) .render(); @@ -7323,7 +7346,7 @@ async fn environment_context_uses_session_shell_when_environment_shell_is_absent "{environment_context}" ); - let primary_environment = turn_context + let primary_environment = step_context .environments .turn_environments .first_mut() @@ -7335,6 +7358,7 @@ async fn environment_context_uses_session_shell_when_environment_shell_is_absent let environment_context = crate::context::EnvironmentContext::from_turn_context( &turn_context, + &step_context, session_shell.as_ref(), ) .render(); @@ -7359,7 +7383,7 @@ async fn build_settings_update_items_emits_environment_item_for_time_changes() { let reference_context_item = previous_context.to_turn_context_item(); let update_items = session - .build_settings_update_items(Some(&reference_context_item), ¤t_context) + .build_settings_update_items_for_test(Some(&reference_context_item), ¤t_context) .await; let environment_update = user_input_texts(&update_items) @@ -7387,7 +7411,7 @@ async fn build_settings_update_items_omits_environment_item_when_disabled() { let reference_context_item = previous_context.to_turn_context_item(); let update_items = session - .build_settings_update_items(Some(&reference_context_item), ¤t_context) + .build_settings_update_items_for_test(Some(&reference_context_item), ¤t_context) .await; let user_texts = user_input_texts(&update_items); @@ -7412,7 +7436,7 @@ async fn build_settings_update_items_emits_realtime_start_when_session_becomes_l current_context.realtime_active = true; let update_items = session - .build_settings_update_items( + .build_settings_update_items_for_test( Some(&previous_context.to_turn_context_item()), ¤t_context, ) @@ -7440,7 +7464,7 @@ async fn build_settings_update_items_emits_realtime_end_when_session_stops_being current_context.realtime_active = false; let update_items = session - .build_settings_update_items( + .build_settings_update_items_for_test( Some(&previous_context.to_turn_context_item()), ¤t_context, ) @@ -7477,7 +7501,7 @@ async fn build_settings_update_items_uses_previous_turn_settings_for_realtime_en .set_previous_turn_settings(Some(previous_turn_settings)) .await; let update_items = session - .build_settings_update_items(Some(&previous_context_item), ¤t_context) + .build_settings_update_items_for_test(Some(&previous_context_item), ¤t_context) .await; let developer_texts = developer_input_texts(&update_items); @@ -7494,7 +7518,7 @@ async fn build_initial_context_uses_previous_realtime_state() { let (session, mut turn_context) = make_session_and_context().await; turn_context.realtime_active = true; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let developer_texts = developer_input_texts(&initial_context); assert!( developer_texts @@ -7508,7 +7532,7 @@ async fn build_initial_context_uses_previous_realtime_state() { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item)); } - let resumed_context = session.build_initial_context(&turn_context).await; + let resumed_context = session.build_initial_context_for_test(&turn_context).await; let resumed_developer_texts = developer_input_texts(&resumed_context); assert!( !resumed_developer_texts @@ -7578,7 +7602,7 @@ async fn build_initial_context_includes_prompt_fragments_from_extensions() { .thread_extension_data .insert(PromptExtensionTestState); - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let developer_messages = developer_message_texts(&initial_context); assert!( @@ -7595,7 +7619,7 @@ async fn build_initial_context_omits_prompt_fragments_without_extension_state() let (mut session, turn_context) = make_session_and_context().await; session.services.extensions = prompt_extension_test_registry(); - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let developer_messages = developer_message_texts(&initial_context); assert!( @@ -7612,7 +7636,9 @@ async fn build_initial_context_adds_multi_agent_v2_root_usage_hint_as_developer_ let (session, turn_context) = make_multi_agent_v2_usage_hint_test_session(/*enable_multi_agent_v2*/ true).await; - let initial_context = session.build_initial_context(turn_context.as_ref()).await; + let initial_context = session + .build_initial_context_for_test(turn_context.as_ref()) + .await; let developer_messages = developer_message_texts(&initial_context); assert!( @@ -7650,7 +7676,9 @@ async fn build_initial_context_adds_multi_agent_v2_subagent_usage_hint_as_develo .expect("thread settings should not be shared") .session_source = session_source; - let initial_context = session.build_initial_context(turn_context.as_ref()).await; + let initial_context = session + .build_initial_context_for_test(turn_context.as_ref()) + .await; let developer_messages = developer_message_texts(&initial_context); assert!( @@ -7672,7 +7700,9 @@ async fn build_initial_context_omits_multi_agent_v2_usage_hints_when_feature_dis let (session, turn_context) = make_multi_agent_v2_usage_hint_test_session(/*enable_multi_agent_v2*/ false).await; - let initial_context = session.build_initial_context(turn_context.as_ref()).await; + let initial_context = session + .build_initial_context_for_test(turn_context.as_ref()) + .await; let developer_messages = developer_message_texts(&initial_context); assert!( @@ -7700,7 +7730,9 @@ async fn build_initial_context_omits_multi_agent_v2_usage_hints_when_hint_disabl ) .await; - let initial_context = session.build_initial_context(turn_context.as_ref()).await; + let initial_context = session + .build_initial_context_for_test(turn_context.as_ref()) + .await; let developer_messages = developer_message_texts(&initial_context); assert!( @@ -7730,7 +7762,7 @@ async fn build_initial_context_omits_default_image_save_location_with_image_hist ) .await; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let developer_texts = developer_input_texts(&initial_context); assert!( !developer_texts @@ -7744,7 +7776,7 @@ async fn build_initial_context_omits_default_image_save_location_with_image_hist async fn build_initial_context_omits_default_image_save_location_without_image_history() { let (session, turn_context) = make_session_and_context().await; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let developer_texts = developer_input_texts(&initial_context); assert!( @@ -7786,7 +7818,7 @@ async fn build_initial_context_trims_skill_metadata_from_context_window_budget() turn_context.model_info.context_window = Some(100); turn_context.turn_skills = TurnSkillsContext::new(HostSkillsSnapshot::new(Arc::new(outcome))); - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let developer_texts = developer_input_texts(&initial_context); assert!( @@ -7937,7 +7969,7 @@ async fn build_initial_context_emits_thread_start_skill_warning_on_repeated_buil turn_context.model_info.context_window = Some(100); turn_context.turn_skills = TurnSkillsContext::new(HostSkillsSnapshot::new(Arc::new(outcome))); - let _ = session.build_initial_context(&turn_context).await; + let _ = session.build_initial_context_for_test(&turn_context).await; let warning_event = timeout(Duration::from_secs(1), rx.recv()) .await .expect("warning event should arrive") @@ -7948,7 +7980,7 @@ async fn build_initial_context_emits_thread_start_skill_warning_on_repeated_buil if message == "Exceeded skills context budget of 2%. All skill descriptions were removed and 2 additional skills were not included in the model-visible skills list." )); - let _ = session.build_initial_context(&turn_context).await; + let _ = session.build_initial_context_for_test(&turn_context).await; let warning_event = timeout(Duration::from_secs(1), rx.recv()) .await .expect("warning event should arrive on repeated build") @@ -7986,7 +8018,7 @@ async fn handle_output_item_done_records_image_save_history_message() { turn_store: Arc::new(codex_extension_api::ExtensionData::new( turn_context.sub_id.clone(), )), - tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)), + tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)).await, cancellation_token: CancellationToken::new(), }; handle_output_item_done(&mut ctx, item.clone(), /*previously_active_item*/ None) @@ -8043,7 +8075,7 @@ async fn handle_output_item_done_skips_image_save_message_when_save_fails() { turn_store: Arc::new(codex_extension_api::ExtensionData::new( turn_context.sub_id.clone(), )), - tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)), + tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)).await, cancellation_token: CancellationToken::new(), }; handle_output_item_done(&mut ctx, item.clone(), /*previously_active_item*/ None) @@ -8068,7 +8100,7 @@ async fn build_initial_context_uses_previous_turn_settings_for_realtime_end() { session .set_previous_turn_settings(Some(previous_turn_settings)) .await; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let developer_texts = developer_input_texts(&initial_context); assert!( developer_texts @@ -8091,7 +8123,7 @@ async fn build_initial_context_restates_realtime_start_when_reference_context_is session .set_previous_turn_settings(Some(previous_turn_settings)) .await; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let developer_texts = developer_input_texts(&initial_context); assert!( developer_texts @@ -8120,15 +8152,7 @@ fn file_system_policy_with_unreadable_glob(turn_context: &TurnContext) -> FileSy #[tokio::test] async fn turn_context_item_stores_local_cwd() { - let (_session, mut turn_context) = make_session_and_context().await; - let environment = turn_context.environments.turn_environments[0].clone(); - let cwd = PathUri::parse("file:///C:/windows").expect("Windows cwd URI"); - turn_context.environments.turn_environments[0] = TurnEnvironment::new( - "remote".to_string(), - environment.environment, - cwd, - environment.shell, - ); + let (_session, turn_context) = make_session_and_context().await; #[allow(deprecated)] let local_cwd = turn_context.cwd.clone(); @@ -8175,10 +8199,10 @@ async fn record_context_updates_and_set_reference_context_item_injects_full_cont { let (session, turn_context) = make_session_and_context().await; session - .record_context_updates_and_set_reference_context_item(&turn_context) + .record_context_updates_and_set_reference_context_item_for_test(&turn_context) .await; let history = session.clone_history().await; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; assert_eq!(history.raw_items().to_vec(), initial_context); let current_context = session.reference_context_item().await; @@ -8206,7 +8230,7 @@ async fn record_context_updates_and_set_reference_context_item_reinjects_full_co .record_conversation_items(&turn_context, std::slice::from_ref(&compacted_summary)) .await; session - .record_context_updates_and_set_reference_context_item(&turn_context) + .record_context_updates_and_set_reference_context_item_for_test(&turn_context) .await; { let mut state = session.state.lock().await; @@ -8220,12 +8244,12 @@ async fn record_context_updates_and_set_reference_context_item_reinjects_full_co .await; session - .record_context_updates_and_set_reference_context_item(&turn_context) + .record_context_updates_and_set_reference_context_item_for_test(&turn_context) .await; let history = session.clone_history().await; let mut expected_history = vec![compacted_summary]; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; expected_history.extend(initial_context); assert_eq!(history.raw_items().to_vec(), expected_history); } @@ -8250,12 +8274,12 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline let rollout_path = attach_thread_persistence(&mut session).await; let update_items = session - .build_settings_update_items(Some(&previous_context_item), &turn_context) + .build_settings_update_items_for_test(Some(&previous_context_item), &turn_context) .await; assert_eq!(update_items, Vec::new()); session - .record_context_updates_and_set_reference_context_item(&turn_context) + .record_context_updates_and_set_reference_context_item_for_test(&turn_context) .await; assert_eq!( @@ -8302,7 +8326,7 @@ async fn record_context_updates_and_set_reference_context_item_persists_split_fi let rollout_path = attach_thread_persistence(&mut session).await; session - .record_context_updates_and_set_reference_context_item(&turn_context) + .record_context_updates_and_set_reference_context_item_for_test(&turn_context) .await; session.ensure_rollout_materialized().await; session.flush_rollout().await.expect("rollout should flush"); @@ -8335,7 +8359,7 @@ async fn build_initial_context_prepends_model_switch_message() { session .set_previous_turn_settings(Some(previous_turn_settings)) .await; - let initial_context = session.build_initial_context(&turn_context).await; + let initial_context = session.build_initial_context_for_test(&turn_context).await; let ResponseItem::Message { role, content, .. } = &initial_context[0] else { panic!("expected developer message"); @@ -8386,7 +8410,7 @@ async fn record_context_updates_and_set_reference_context_item_persists_full_rei })) .await; session - .record_context_updates_and_set_reference_context_item(&turn_context) + .record_context_updates_and_set_reference_context_item_for_test(&turn_context) .await; session.ensure_rollout_materialized().await; session.flush_rollout().await.expect("rollout should flush"); @@ -9497,7 +9521,7 @@ async fn tool_calls_reopen_mailbox_delivery_for_current_turn() { sess: Arc::clone(&sess), turn_context: Arc::clone(&tc), turn_store: Arc::new(codex_extension_api::ExtensionData::new(tc.sub_id.clone())), - tool_runtime: test_tool_runtime(Arc::clone(&sess), Arc::clone(&tc)), + tool_runtime: test_tool_runtime(Arc::clone(&sess), Arc::clone(&tc)).await, cancellation_token: CancellationToken::new(), }; @@ -9523,7 +9547,8 @@ async fn abort_review_task_emits_exited_then_aborted_and_records_history() { }], client_id: None, }]; - sess.spawn_task(Arc::clone(&tc), input, ReviewTask::new()) + let environments = sess.services.turn_environments.snapshot().await; + sess.spawn_task(Arc::clone(&tc), input, ReviewTask::new(environments)) .await; sess.abort_all_tasks(TurnAbortReason::Interrupted).await; @@ -9602,8 +9627,10 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { .await }; let deferred_mcp_tools = Some(tools.clone()); - let router = ToolRouter::from_turn_context( + let step_context = step_context_for_session(&session).await; + let router = ToolRouter::from_contexts( &turn_context, + step_context, crate::tools::router::ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools, @@ -9661,7 +9688,7 @@ async fn sample_rollout( // personality_spec) matches reconstruction. let reconstruction_turn = session.new_default_turn().await; let mut initial_context = session - .build_initial_context(reconstruction_turn.as_ref()) + .build_initial_context_for_test(reconstruction_turn.as_ref()) .await; // Ensure personality_spec is present when Personality is enabled, so expected matches // what reconstruction produces (build_initial_context may omit it when baked into model). @@ -9836,6 +9863,7 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() { .expect("test setup should allow updating approval policy"); let session = Arc::new(session); let mut turn_context = Arc::new(turn_context_raw); + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let command_script = "echo hi"; let timeout_ms = 1000; @@ -9853,6 +9881,7 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() { .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step: step_context, cancellation_token: CancellationToken::new(), tracker: Arc::clone(&turn_diff_tracker), call_id, @@ -9929,7 +9958,6 @@ async fn shell_tool_cancellation_waits_for_runtime_cleanup() -> anyhow::Result<( .await?; let turn_context = session.new_default_turn().await; let session = Arc::new(session); - let turn_context = Arc::new(turn_context); let temp_dir = tempfile::TempDir::new()?; let ready_marker = temp_dir.path().join("ready"); let cleanup_marker = temp_dir.path().join("cleanup"); @@ -9959,6 +9987,7 @@ while :; do sleep 1; done"#, let cancellation_tx = cancellation_token.clone(); let handle = tokio::spawn( test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)) + .await .handle_tool_call(call, cancellation_token), ); @@ -9999,6 +10028,7 @@ async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request() .expect("test setup should allow updating approval policy"); let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let handler = ExecCommandHandler::default(); @@ -10006,6 +10036,7 @@ async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request() .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step: step_context, cancellation_token: CancellationToken::new(), tracker: Arc::clone(&tracker), call_id: "exec-call".to_string(), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index b9727a04014..dc033a40b5f 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -116,6 +116,7 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { ); let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let requested_permissions = RequestPermissionProfile { network: Some(NetworkPermissions { @@ -123,7 +124,7 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { }), ..RequestPermissionProfile::default() }; - let environment = turn_context + let environment = step_context .environments .primary() .expect("primary environment") @@ -132,6 +133,7 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { Duration::from_secs(45), session.request_permissions_for_environment( &turn_context, + &step_context.environments, "perm-call-1".to_string(), RequestPermissionsArgs { environment_id: None, @@ -213,13 +215,15 @@ async fn request_permissions_guardian_review_stops_when_cancelled() { ..RequestPermissionProfile::default() }; let cancellation_token = CancellationToken::new(); + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let request_handle = tokio::spawn({ let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); + let step_context = Arc::clone(&step_context); let requested_permissions = requested_permissions.clone(); let cancellation_token = cancellation_token.clone(); async move { - let environment = turn_context + let environment = step_context .environments .primary() .expect("primary environment") @@ -227,6 +231,7 @@ async fn request_permissions_guardian_review_stops_when_cancelled() { session .request_permissions_for_environment( &turn_context, + &step_context.environments, "perm-call-cancelled".to_string(), RequestPermissionsArgs { environment_id: None, @@ -322,6 +327,7 @@ async fn guardian_allows_shell_command_additional_permissions_requests_past_poli ); let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let expiration_ms: u64 = if cfg!(windows) { 2_500 } else { 1_000 }; let handler = crate::tools::handlers::ShellCommandHandler::from( @@ -333,6 +339,7 @@ async fn guardian_allows_shell_command_additional_permissions_requests_past_poli .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step: step_context, cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "test-call".to_string(), @@ -427,6 +434,7 @@ async fn strict_auto_review_turn_grant_forces_guardian_for_shell_command_policy_ ); let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let handler = crate::tools::handlers::ShellCommandHandler::from( codex_tools::ShellCommandBackendConfig::Classic, @@ -437,6 +445,7 @@ async fn strict_auto_review_turn_grant_forces_guardian_for_shell_command_policy_ .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step: step_context, cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "strict-shell-command-call".to_string(), @@ -477,6 +486,7 @@ async fn guardian_allows_unified_exec_additional_permissions_requests_past_polic .expect("test setup should allow enabling request permissions"); let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let handler = ExecCommandHandler::default(); @@ -484,6 +494,7 @@ async fn guardian_allows_unified_exec_additional_permissions_requests_past_polic .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step: step_context, cancellation_token: CancellationToken::new(), tracker: Arc::clone(&tracker), call_id: "exec-call".to_string(), @@ -524,7 +535,7 @@ async fn process_compacted_history_preserves_separate_guardian_developer_message turn_context.session_source = guardian_source; turn_context.developer_instructions = Some(guardian_policy.clone()); - let refreshed = crate::compact_remote::process_compacted_history( + let refreshed = crate::compact_remote::process_compacted_history_for_test( &session, &turn_context, vec![ @@ -600,6 +611,7 @@ async fn shell_command_allows_sticky_turn_permissions_without_inline_request_per let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let handler = crate::tools::handlers::ShellCommandHandler::from( codex_tools::ShellCommandBackendConfig::Classic, @@ -610,6 +622,7 @@ async fn shell_command_allows_sticky_turn_permissions_without_inline_request_per .handle(ToolInvocation { session: Arc::clone(&session), turn: Arc::clone(&turn_context), + step: step_context, cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "sticky-turn-grant".to_string(), diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 4beef792bb3..91a422a3baa 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -42,6 +42,7 @@ use crate::responses_retry::handle_retryable_response_stream_error; use crate::session::PreviousTurnSettings; use crate::session::TurnInput; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::HandleOutputCtx; use crate::stream_events_utils::TurnItemContributorPolicy; @@ -145,13 +146,16 @@ pub(crate) async fn run_turn( prewarmed_client_session: Option, cancellation_token: CancellationToken, ) -> Option { + let step_context = sess.prepare_step_for_request().await; let mut client_session = prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); // TODO(ccunningham): Pre-turn compaction runs before context updates and the // new user message are recorded. Estimate pending incoming items (context // diffs/full reinjection + user input) and trigger compaction preemptively // when they would push the thread over the compaction threshold. - if let Err(err) = run_pre_sampling_compact(&sess, &turn_context, &mut client_session).await { + if let Err(err) = + run_pre_sampling_compact(&sess, &turn_context, &step_context, &mut client_session).await + { let error = err.to_codex_protocol_error(); sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) .await; @@ -159,11 +163,20 @@ pub(crate) async fn run_turn( return None; } - sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref()) - .await; + sess.record_context_updates_and_set_reference_context_item( + turn_context.as_ref(), + step_context.as_ref(), + ) + .await; - let (injection_items, explicitly_enabled_connectors) = - build_skills_and_plugins(&sess, turn_context.as_ref(), &input, &cancellation_token).await?; + let (injection_items, explicitly_enabled_connectors) = build_skills_and_plugins( + &sess, + turn_context.as_ref(), + step_context.as_ref(), + &input, + &cancellation_token, + ) + .await?; if run_pending_session_start_hooks(&sess, &turn_context).await { return None; @@ -192,7 +205,7 @@ pub(crate) async fn run_turn( let mut stop_hook_active = false; // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains // many turns, from the perspective of the user, it is a single turn. - let display_roots = turn_diff_display_roots(turn_context.as_ref()).await; + let display_roots = turn_diff_display_roots(step_context.as_ref()).await; let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new( TurnDiffTracker::with_environment_display_roots(display_roots), )); @@ -237,6 +250,7 @@ pub(crate) async fn run_turn( match run_sampling_request( Arc::clone(&sess), Arc::clone(&turn_context), + Arc::clone(&step_context), Arc::clone(&turn_extension_data), Arc::clone(&turn_diff_tracker), &mut client_session, @@ -293,7 +307,7 @@ pub(crate) async fn run_turn( .await; let started_new_context_window = sess - .maybe_start_new_context_window(turn_context.as_ref()) + .maybe_start_new_context_window(turn_context.as_ref(), step_context.as_ref()) .await .is_some(); if started_new_context_window && needs_follow_up { @@ -306,6 +320,7 @@ pub(crate) async fn run_turn( if let Err(err) = run_auto_compact( &sess, &turn_context, + &step_context, &mut client_session, InitialContextInjection::BeforeLastUserMessage, CompactionReason::ContextLimit, @@ -414,9 +429,9 @@ pub(crate) async fn run_turn( } #[instrument(level = "trace", skip_all)] -async fn turn_diff_display_roots(turn_context: &TurnContext) -> Vec<(String, PathBuf)> { +async fn turn_diff_display_roots(step_context: &StepContext) -> Vec<(String, PathBuf)> { let mut display_roots = Vec::new(); - for turn_environment in &turn_context.environments.turn_environments { + for turn_environment in &step_context.environments.turn_environments { // TODO(anp): Migrate git-root discovery and diff display roots to PathUri so foreign // environment roots can participate without host-native conversion. let Ok(cwd) = turn_environment.cwd().to_abs_path() else { @@ -465,6 +480,7 @@ async fn run_hooks_and_record_inputs( async fn build_skills_and_plugins( sess: &Arc, turn_context: &TurnContext, + step_context: &StepContext, input: &[TurnInput], cancellation_token: &CancellationToken, ) -> Option<(Vec, HashSet)> { @@ -530,9 +546,14 @@ 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) - .await?; + let extension_injection_items = build_extension_turn_input_items( + sess, + turn_context, + step_context, + &user_input, + cancellation_token, + ) + .await?; let skill_name_counts_lower = build_skill_name_counts(&skills_outcome.skills, &skills_outcome.disabled_paths).1; let mentioned_skills = collect_explicit_skill_mentions( @@ -627,6 +648,7 @@ async fn build_skills_and_plugins( async fn build_extension_turn_input_items( sess: &Arc, turn_context: &TurnContext, + step_context: &StepContext, user_input: &[UserInput], cancellation_token: &CancellationToken, ) -> Option> { @@ -635,7 +657,7 @@ async fn build_extension_turn_input_items( return Some(Vec::new()); } - let environments = turn_context + let environments = step_context .environments .turn_environments .iter() @@ -799,15 +821,18 @@ async fn auto_compact_token_status( async fn run_pre_sampling_compact( sess: &Arc, turn_context: &Arc, + step_context: &Arc, client_session: &mut ModelClientSession, ) -> CodexResult<()> { - maybe_run_previous_model_inline_compact(sess, turn_context, client_session).await?; + maybe_run_previous_model_inline_compact(sess, turn_context, step_context, client_session) + .await?; let token_status = auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await; // Compact if the configured auto-compaction budget or usable context window is exhausted. if token_status.token_limit_reached { run_auto_compact( sess, turn_context, + step_context, client_session, InitialContextInjection::DoNotInject, CompactionReason::ContextLimit, @@ -833,6 +858,7 @@ fn comp_hash_changed(previous: Option<&str>, current: Option<&str>) -> bool { async fn maybe_run_previous_model_inline_compact( sess: &Arc, turn_context: &Arc, + step_context: &Arc, client_session: &mut ModelClientSession, ) -> CodexResult<()> { let Some(previous_turn_settings) = sess.previous_turn_settings().await else { @@ -852,6 +878,7 @@ async fn maybe_run_previous_model_inline_compact( run_auto_compact( sess, &previous_model_turn_context, + step_context, client_session, InitialContextInjection::DoNotInject, CompactionReason::CompHashChanged, @@ -889,6 +916,7 @@ async fn maybe_run_previous_model_inline_compact( run_auto_compact( sess, &previous_model_turn_context, + step_context, client_session, InitialContextInjection::DoNotInject, CompactionReason::ModelDownshift, @@ -907,6 +935,7 @@ async fn maybe_run_previous_model_inline_compact( async fn run_auto_compact( sess: &Arc, turn_context: &Arc, + step_context: &Arc, client_session: &mut ModelClientSession, initial_context_injection: InitialContextInjection, reason: CompactionReason, @@ -926,6 +955,7 @@ async fn run_auto_compact( run_inline_remote_auto_compact_task_v2( Arc::clone(sess), Arc::clone(turn_context), + Arc::clone(step_context), client_session, initial_context_injection, reason, @@ -942,6 +972,7 @@ async fn run_auto_compact( run_inline_remote_auto_compact_task( Arc::clone(sess), Arc::clone(turn_context), + Arc::clone(step_context), client_session.turn_state(), initial_context_injection, reason, @@ -957,6 +988,7 @@ async fn run_auto_compact( run_inline_auto_compact_task( Arc::clone(sess), Arc::clone(turn_context), + Arc::clone(step_context), initial_context_injection, reason, phase, @@ -1049,6 +1081,7 @@ pub(crate) fn build_prompt( async fn run_sampling_request( sess: Arc, turn_context: Arc, + step_context: Arc, turn_store: Arc, turn_diff_tracker: SharedTurnDiffTracker, client_session: &mut ModelClientSession, @@ -1056,7 +1089,13 @@ async fn run_sampling_request( input: Vec, cancellation_token: CancellationToken, ) -> CodexResult<(SamplingRequestResult, Vec)> { - let router = built_tools(sess.as_ref(), turn_context.as_ref(), &cancellation_token).await?; + let router = built_tools( + sess.as_ref(), + turn_context.as_ref(), + Arc::clone(&step_context), + &cancellation_token, + ) + .await?; let base_instructions = sess.get_base_instructions().await; @@ -1153,6 +1192,7 @@ async fn run_sampling_request( pub(crate) async fn built_tools( sess: &Session, turn_context: &TurnContext, + step_context: Arc, cancellation_token: &CancellationToken, ) -> CodexResult> { let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); @@ -1270,8 +1310,9 @@ pub(crate) async fn built_tools( ); let mcp_tools = has_mcp_servers.then_some(mcp_tool_exposure.direct_tools); let deferred_mcp_tools = mcp_tool_exposure.deferred_tools; - Ok(Arc::new(ToolRouter::from_turn_context( + Ok(Arc::new(ToolRouter::from_contexts( turn_context, + step_context, ToolRouterParams { mcp_tools, deferred_mcp_tools, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 91880bdefa0..307fb20a7b8 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -1,6 +1,5 @@ use super::*; use crate::agents_md::LoadedAgentsMd; -use crate::environment_selection::TurnEnvironmentSnapshot; use crate::shell_snapshot::ShellSnapshotFile; use codex_core_skills::HostSkillsSnapshot; use codex_file_system::FileSystemSandboxContext; @@ -113,7 +112,6 @@ pub struct TurnContext { pub(crate) reasoning_summary: ReasoningSummaryConfig, pub(crate) session_source: SessionSource, pub(crate) parent_thread_id: Option, - pub(crate) environments: TurnEnvironmentSnapshot, /// The session's absolute working directory. All relative paths provided /// by the model as well as sandbox policies are resolved against this path /// instead of `std::env::current_dir()`. @@ -205,10 +203,6 @@ impl TurnContext { .apps_enabled_for_auth(uses_codex_backend) } - pub(crate) fn tool_environment_mode(&self) -> ToolEnvironmentMode { - ToolEnvironmentMode::from_count(self.environments.turn_environments.len()) - } - pub(crate) async fn with_model( &self, model: String, @@ -267,7 +261,6 @@ impl TurnContext { reasoning_summary: self.reasoning_summary, session_source: self.session_source.clone(), parent_thread_id: self.parent_thread_id, - environments: self.environments.clone(), #[allow(deprecated)] cwd: self.cwd.clone(), current_date: self.current_date.clone(), @@ -300,13 +293,6 @@ impl TurnContext { } } - #[deprecated(note = "resolve paths from the selected turn environment cwd instead")] - pub(crate) fn resolve_path(&self, path: Option) -> AbsolutePathBuf { - #[allow(deprecated)] - path.as_ref() - .map_or_else(|| self.cwd.clone(), |path| self.cwd.join(path)) - } - pub(crate) fn file_system_sandbox_context( &self, additional_permissions: Option, @@ -490,7 +476,6 @@ impl Session { model_info: ModelInfo, models_manager: &SharedModelsManager, network: Option, - environments: TurnEnvironmentSnapshot, cwd: AbsolutePathBuf, sub_id: String, skills_snapshot: HostSkillsSnapshot, @@ -550,7 +535,6 @@ impl Session { reasoning_summary, session_source, parent_thread_id: session_configuration.parent_thread_id, - environments, #[allow(deprecated)] cwd, current_date: Some(current_date), @@ -760,7 +744,6 @@ impl Session { ) .then(|| started_proxy.proxy()) }), - turn_environments, cwd, sub_id, skills_snapshot, @@ -770,15 +753,10 @@ impl Session { if let Some(final_schema) = final_output_json_schema { turn_context.final_output_json_schema = final_schema; } - let turn_context = Arc::new(turn_context); - if turn_context - .environments - .single_local_environment_cwd() - .is_some() - { + if turn_environments.single_local_environment_cwd().is_some() { turn_context.turn_metadata_state.spawn_git_enrichment_task(); } - turn_context + Arc::new(turn_context) } pub(crate) async fn maybe_emit_unknown_model_warning_for_turn(&self, tc: &TurnContext) { diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 0474eb87046..79fbc952d64 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -243,10 +243,12 @@ async fn schedule_startup_prewarm_inner( /*status*/ None, ); let startup_cancellation_token = CancellationToken::new(); + let step_context = session.prepare_step_for_request().await; let built_tools_started_at = Instant::now(); let startup_router = built_tools( session.as_ref(), startup_turn_context.as_ref(), + step_context, &startup_cancellation_token, ) .await?; diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index b89d4f5c6e5..e15190f3c7d 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -8,6 +8,7 @@ use super::image_generation_artifact_path; use super::last_assistant_message_from_item; use super::response_item_may_include_external_context; use super::save_image_generation_result; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::ToolRouter; use crate::tools::parallel::ToolCallRuntime; @@ -276,8 +277,10 @@ async fn handle_output_item_done_returns_contributed_last_agent_message() { session.services.extensions = Arc::new(builder.build()); let session = Arc::new(session); let turn_context = Arc::new(turn_context); - let router = Arc::new(ToolRouter::from_turn_context( + let step_context = Arc::new(StepContext::local_for_test(&turn_context)); + let router = Arc::new(ToolRouter::from_contexts( &turn_context, + step_context, crate::tools::router::ToolRouterParams { tool_suggest_candidates: None, mcp_tools: None, diff --git a/codex-rs/core/src/tasks/compact.rs b/codex-rs/core/src/tasks/compact.rs index dd9d6cb079e..cb3b7aea02c 100644 --- a/codex-rs/core/src/tasks/compact.rs +++ b/codex-rs/core/src/tasks/compact.rs @@ -29,6 +29,7 @@ impl SessionTask for CompactTask { _cancellation_token: CancellationToken, ) -> Option { let session = session.clone_session(); + let step = session.prepare_step_for_request().await; let _ = if crate::compact::should_use_remote_compact_task(ctx.provider.info()) { if ctx .config @@ -40,14 +41,14 @@ impl SessionTask for CompactTask { "remote_v2", /*manual*/ true, ); - crate::compact_remote_v2::run_remote_compact_task(session.clone(), ctx).await + crate::compact_remote_v2::run_remote_compact_task(session.clone(), ctx, step).await } else { emit_compact_metric( &session.services.session_telemetry, "remote", /*manual*/ true, ); - crate::compact_remote::run_remote_compact_task(session.clone(), ctx).await + crate::compact_remote::run_remote_compact_task(session.clone(), ctx, step).await } } else { emit_compact_metric( @@ -65,7 +66,7 @@ impl SessionTask for CompactTask { // Compaction prompt is synthesized; no UI element ranges to preserve. text_elements: Vec::new(), }]; - crate::compact::run_compact_task(session.clone(), ctx, input).await + crate::compact::run_compact_task(session.clone(), ctx, step, input).await }; None } diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index bb8670237e2..9abae191eed 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -18,6 +18,7 @@ use tokio_util::sync::CancellationToken; use crate::codex_delegate::run_codex_thread_one_shot; use crate::config::Constrained; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::review_format::format_review_findings_block; use crate::review_format::render_review_output_text; use crate::session::TurnInput; @@ -30,12 +31,14 @@ use codex_protocol::user_input::UserInput; use super::SessionTask; use super::SessionTaskContext; -#[derive(Clone, Copy)] -pub(crate) struct ReviewTask; +#[derive(Clone)] +pub(crate) struct ReviewTask { + environments: TurnEnvironmentSnapshot, +} impl ReviewTask { - pub(crate) fn new() -> Self { - Self + pub(crate) fn new(environments: TurnEnvironmentSnapshot) -> Self { + Self { environments } } } @@ -73,6 +76,7 @@ impl SessionTask for ReviewTask { let output = match start_review_conversation( session.clone(), ctx.clone(), + self.environments.clone(), user_input, cancellation_token.clone(), ) @@ -95,6 +99,7 @@ impl SessionTask for ReviewTask { async fn start_review_conversation( session: Arc, ctx: Arc, + environments: TurnEnvironmentSnapshot, input: Vec, cancellation_token: CancellationToken, ) -> Option> { @@ -128,6 +133,7 @@ async fn start_review_conversation( input, session.clone_session(), ctx.clone(), + environments, cancellation_token, SubAgentSource::Review, /*final_output_json_schema*/ None, diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index 057718b9f7c..c992efee8e1 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -125,8 +125,8 @@ pub(crate) async fn execute_user_shell_command( session.send_event(turn_context.as_ref(), event).await; } - let Some((turn_environment, environment_shell)) = turn_context - .environments + let environments = session.services.turn_environments.snapshot().await; + let Some((turn_environment, environment_shell)) = environments .local() .and_then(|environment| environment.shell.as_ref().map(|shell| (environment, shell))) else { diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index cec54b5e444..17afdf5a58f 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -225,7 +225,7 @@ fn out_of_range_truncation_drops_pre_user_active_turn_prefix() { #[tokio::test] async fn ignores_session_prefix_messages_when_truncating() { let (session, turn_context) = make_session_and_context().await; - let mut items = session.build_initial_context(&turn_context).await; + let mut items = session.build_initial_context_for_test(&turn_context).await; items.push(user_msg("feature request")); items.push(assistant_msg("ack")); items.push(user_msg("second question")); @@ -603,15 +603,24 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { .new_turn_with_sub_id("resume-turn".to_string(), SessionSettingsUpdate::default()) .await .expect("build resumed turn context"); - assert_eq!(resumed_turn.environments.turn_environments.len(), 1); + let resumed_environments = resumed + .thread + .codex + .session + .services + .turn_environments + .snapshot() + .await; + assert_eq!(resumed_environments.turn_environments.len(), 1); assert_eq!( - resumed_turn.environments.turn_environments[0].cwd(), + resumed_environments.turn_environments[0].cwd(), &PathUri::from_abs_path(&default_cwd) ); assert_ne!( - resumed_turn.environments.turn_environments[0].cwd(), + resumed_environments.turn_environments[0].cwd(), &PathUri::from_abs_path(&selected_cwd) ); + assert_eq!(resumed_turn.config.cwd.as_path(), default_cwd.as_path()); let forked = manager .fork_thread( @@ -630,15 +639,24 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { .new_turn_with_sub_id("fork-turn".to_string(), SessionSettingsUpdate::default()) .await .expect("build forked turn context"); - assert_eq!(forked_turn.environments.turn_environments.len(), 1); + let forked_environments = forked + .thread + .codex + .session + .services + .turn_environments + .snapshot() + .await; + assert_eq!(forked_environments.turn_environments.len(), 1); assert_eq!( - forked_turn.environments.turn_environments[0].cwd(), + forked_environments.turn_environments[0].cwd(), &PathUri::from_abs_path(&default_cwd) ); assert_ne!( - forked_turn.environments.turn_environments[0].cwd(), + forked_environments.turn_environments[0].cwd(), &PathUri::from_abs_path(&selected_cwd) ); + assert_eq!(forked_turn.config.cwd.as_path(), default_cwd.as_path()); } #[tokio::test] diff --git a/codex-rs/core/src/thread_rollout_truncation_tests.rs b/codex-rs/core/src/thread_rollout_truncation_tests.rs index bb6c3ac03e7..2ec6b9e21d9 100644 --- a/codex-rs/core/src/thread_rollout_truncation_tests.rs +++ b/codex-rs/core/src/thread_rollout_truncation_tests.rs @@ -166,7 +166,7 @@ fn truncates_rollout_from_start_applies_thread_rollback_markers() { #[tokio::test] async fn ignores_session_prefix_messages_when_truncating_rollout_from_start() { let (session, turn_context) = make_session_and_context().await; - let mut items = session.build_initial_context(&turn_context).await; + let mut items = session.build_initial_context_for_test(&turn_context).await; items.push(user_msg("feature request")); items.push(assistant_msg("ack")); items.push(user_msg("second question")); diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 15177944ca0..677a37aff05 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -14,17 +14,20 @@ use tokio::sync::oneshot; use tokio::sync::watch; use tokio_util::sync::CancellationToken; -use super::ExecContext; use super::PUBLIC_TOOL_NAME; use super::call_nested_tool; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::parallel::ToolCallRuntime; +type CellHostSender = watch::Sender>>; +type CellHostMap = Mutex>; + pub(super) struct CodeModeDispatchBroker { dispatch_tx: async_channel::Sender, dispatch_rx: async_channel::Receiver, - dispatch_gates: Arc>>>, + dispatch_hosts: Arc, + current_host: Arc>>>, } impl CodeModeDispatchBroker { @@ -33,33 +36,43 @@ impl CodeModeDispatchBroker { Self { dispatch_tx, dispatch_rx, - dispatch_gates: Arc::new(Mutex::new(HashMap::new())), + dispatch_hosts: Arc::new(Mutex::new(HashMap::new())), + current_host: Arc::new(Mutex::new(None)), } } - pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) { - dispatch_gate(&self.dispatch_gates, cell_id).send_replace(true); + pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) -> Result<(), String> { + let host = match self.current_host.lock() { + Ok(current_host) => current_host.clone(), + Err(poisoned) => poisoned.into_inner().clone(), + } + .ok_or_else(|| "code mode tool dispatcher is unavailable".to_string())?; + dispatch_host(&self.dispatch_hosts, cell_id).send_replace(Some(host)); + Ok(()) } pub(super) fn close_cell(&self, cell_id: &CellId) { - remove_dispatch_gate(&self.dispatch_gates, cell_id); + remove_dispatch_host(&self.dispatch_hosts, cell_id); } pub(super) fn start_turn_worker( &self, - exec: ExecContext, + session: Arc, + turn: Arc, router: Arc, tracker: SharedTurnDiffTracker, ) -> CodeModeDispatchWorker { - let tool_runtime = ToolCallRuntime::new( - router, - Arc::clone(&exec.session), - Arc::clone(&exec.turn), - tracker, - ); - let host = Arc::new(CoreTurnHost { exec, tool_runtime }); + let tool_runtime = ToolCallRuntime::new(router, Arc::clone(&session), turn, tracker); + let host = Arc::new(CoreTurnHost { + session, + tool_runtime, + }); + match self.current_host.lock() { + Ok(mut current_host) => *current_host = Some(Arc::clone(&host)), + Err(poisoned) => *poisoned.into_inner() = Some(Arc::clone(&host)), + } let dispatch_rx = self.dispatch_rx.clone(); - let dispatch_gates = Arc::clone(&self.dispatch_gates); + let dispatch_hosts = Arc::clone(&self.dispatch_hosts); let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); tokio::spawn(async move { loop { @@ -78,16 +91,12 @@ impl CodeModeDispatchBroker { cancellation_token, response_tx, } => { - let response = if wait_until_cell_ready_for_dispatch( - &dispatch_gates, - &cell_id, - &cancellation_token, - ) - .await + let response = if let Some(host) = + wait_for_cell_host(&dispatch_hosts, &cell_id, &cancellation_token).await { host.notify(call_id, cell_id, text).await } else { - remove_dispatch_gate(&dispatch_gates, &cell_id); + remove_dispatch_host(&dispatch_hosts, &cell_id); Err("code mode notification cancelled".to_string()) }; let _ = response_tx.send(response); @@ -98,17 +107,13 @@ impl CodeModeDispatchBroker { response_tx, } => { let cell_id = invocation.cell_id.clone(); - if !wait_until_cell_ready_for_dispatch( - &dispatch_gates, - &cell_id, - &cancellation_token, - ) - .await - { - remove_dispatch_gate(&dispatch_gates, &cell_id); + let Some(host) = + wait_for_cell_host(&dispatch_hosts, &cell_id, &cancellation_token) + .await + else { + remove_dispatch_host(&dispatch_hosts, &cell_id); continue; - } - let host = Arc::clone(&host); + }; tokio::spawn(async move { let response = tokio::select! { response = host.invoke_tool( @@ -125,55 +130,51 @@ impl CodeModeDispatchBroker { }); CodeModeDispatchWorker { shutdown_tx: Some(shutdown_tx), + current_host: Arc::clone(&self.current_host), + host, } } } -fn dispatch_gate( - dispatch_gates: &Mutex>>, - cell_id: &CellId, -) -> watch::Sender { - let mut dispatch_gates = match dispatch_gates.lock() { - Ok(dispatch_gates) => dispatch_gates, +fn dispatch_host(dispatch_hosts: &CellHostMap, cell_id: &CellId) -> CellHostSender { + let mut dispatch_hosts = match dispatch_hosts.lock() { + Ok(dispatch_hosts) => dispatch_hosts, Err(poisoned) => poisoned.into_inner(), }; - dispatch_gates + dispatch_hosts .entry(cell_id.clone()) - .or_insert_with(|| watch::channel(false).0) + .or_insert_with(|| watch::channel(None).0) .clone() } -fn remove_dispatch_gate( - dispatch_gates: &Mutex>>, - cell_id: &CellId, -) { - let mut dispatch_gates = match dispatch_gates.lock() { - Ok(dispatch_gates) => dispatch_gates, +fn remove_dispatch_host(dispatch_hosts: &CellHostMap, cell_id: &CellId) { + let mut dispatch_hosts = match dispatch_hosts.lock() { + Ok(dispatch_hosts) => dispatch_hosts, Err(poisoned) => poisoned.into_inner(), }; - dispatch_gates.remove(cell_id); + dispatch_hosts.remove(cell_id); } -async fn wait_until_cell_ready_for_dispatch( - dispatch_gates: &Mutex>>, +async fn wait_for_cell_host( + dispatch_hosts: &CellHostMap, cell_id: &CellId, cancellation_token: &CancellationToken, -) -> bool { +) -> Option> { if cancellation_token.is_cancelled() { - return false; + return None; } - let mut ready_rx = dispatch_gate(dispatch_gates, cell_id).subscribe(); + let mut host_rx = dispatch_host(dispatch_hosts, cell_id).subscribe(); loop { - if *ready_rx.borrow_and_update() { - return true; + if let Some(host) = host_rx.borrow_and_update().clone() { + return Some(host); } tokio::select! { - changed = ready_rx.changed() => { + changed = host_rx.changed() => { if changed.is_err() { - return false; + return None; } } - _ = cancellation_token.cancelled() => return false, + _ = cancellation_token.cancelled() => return None, } } } @@ -261,10 +262,22 @@ enum DispatchMessage { pub(crate) struct CodeModeDispatchWorker { shutdown_tx: Option>, + current_host: Arc>>>, + host: Arc, } impl Drop for CodeModeDispatchWorker { fn drop(&mut self) { + let mut current_host = match self.current_host.lock() { + Ok(current_host) => current_host, + Err(poisoned) => poisoned.into_inner(), + }; + if current_host + .as_ref() + .is_some_and(|current_host| Arc::ptr_eq(current_host, &self.host)) + { + *current_host = None; + } if let Some(shutdown_tx) = self.shutdown_tx.take() { let _ = shutdown_tx.send(()); } @@ -272,7 +285,7 @@ impl Drop for CodeModeDispatchWorker { } struct CoreTurnHost { - exec: ExecContext, + session: Arc, tool_runtime: ToolCallRuntime, } @@ -282,22 +295,16 @@ impl CoreTurnHost { invocation: CodeModeNestedToolCall, cancellation_token: CancellationToken, ) -> Result { - call_nested_tool( - self.exec.clone(), - self.tool_runtime.clone(), - invocation, - cancellation_token, - ) - .await - .map_err(|error| error.to_string()) + call_nested_tool(self.tool_runtime.clone(), invocation, cancellation_token) + .await + .map_err(|error| error.to_string()) } async fn notify(&self, call_id: String, cell_id: CellId, text: String) -> Result<(), String> { if text.trim().is_empty() { return Ok(()); } - self.exec - .session + self.session .inject_if_running(vec![ResponseItem::CustomToolCallOutput { call_id, name: Some(PUBLIC_TOOL_NAME.to_string()), diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 2e46c1bd6b9..0b169f985d3 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -67,7 +67,8 @@ impl CodeModeExecuteHandler { exec.session .services .code_mode_service - .mark_cell_ready_for_dispatch(&cell_id); + .mark_cell_ready_for_dispatch(&cell_id) + .map_err(FunctionCallError::RespondToModel)?; let response = started_cell .initial_response() .await diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 09918033e66..daa416cb856 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -102,8 +102,11 @@ impl CodeModeService { } } - pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) { - self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id); + pub(crate) fn mark_cell_ready_for_dispatch( + &self, + cell_id: &codex_code_mode::CellId, + ) -> Result<(), String> { + self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id) } pub(crate) fn finish_cell_dispatch(&self, cell_id: &CellId) { @@ -124,14 +127,12 @@ impl CodeModeService { return None; } - let exec = ExecContext { - session: Arc::clone(session), - turn: Arc::clone(turn), - }; - Some( - self.dispatch_broker - .start_turn_worker(exec, router, tracker), - ) + Some(self.dispatch_broker.start_turn_worker( + Arc::clone(session), + Arc::clone(turn), + router, + tracker, + )) } fn session(&self) -> Result<&Arc, String> { @@ -236,7 +237,6 @@ fn truncate_code_mode_result( } async fn call_nested_tool( - _exec: ExecContext, tool_runtime: ToolCallRuntime, invocation: CodeModeNestedToolCall, cancellation_token: CancellationToken, diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 57cf0d80b37..08fd00a9d87 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -1,6 +1,7 @@ use crate::context_manager::truncate_function_output_payload; use crate::original_image_detail::sanitize_original_image_detail; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::TELEMETRY_PREVIEW_MAX_BYTES; use crate::tools::TELEMETRY_PREVIEW_MAX_LINES; @@ -54,6 +55,7 @@ pub enum ToolCallSource { pub struct ToolInvocation { pub session: Arc, pub turn: Arc, + pub step: Arc, pub cancellation_token: CancellationToken, pub tracker: SharedTurnDiffTracker, pub call_id: String, diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index 360744973c5..1c9a4c7a599 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -12,6 +12,7 @@ use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; use codex_utils_absolute_path::AbsolutePathBuf; use futures::StreamExt; @@ -155,7 +156,7 @@ fn normalize_max_runtime_seconds(requested: Option) -> Result, async fn run_agent_job_loop( session: Arc, - turn: Arc, + environments: Vec, db: Arc, job_id: String, options: JobRunnerOptions, @@ -209,7 +210,7 @@ async fn run_agent_job_loop( )))), SpawnAgentOptions { parent_thread_id: Some(session.thread_id), - environments: Some(turn.environments.to_selections()), + environments: Some(environments.clone()), ..Default::default() }, ) diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs index a519e5db1b0..c53e9a06e8e 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs @@ -1,4 +1,6 @@ +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::function_tool::FunctionCallError; +use crate::session::step_context::StepContext; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; @@ -36,6 +38,7 @@ impl SpawnAgentsOnCsvHandler { let ToolInvocation { session, turn, + step, payload, .. } = invocation; @@ -49,7 +52,7 @@ impl SpawnAgentsOnCsvHandler { } }; - handle(session, turn, arguments) + handle(session, turn, step, arguments) .await .map(boxed_tool_output) } @@ -69,6 +72,7 @@ impl CoreToolRuntime for SpawnAgentsOnCsvHandler { pub async fn handle( session: Arc, turn: Arc, + step: Arc, arguments: String, ) -> Result { let args: SpawnAgentsOnCsvArgs = parse_arguments(arguments.as_str())?; @@ -78,7 +82,8 @@ pub async fn handle( )); } - let cwd = single_local_environment_cwd(&turn)?; + let environments = step.environments.clone(); + let cwd = single_local_environment_cwd(&environments)?; let db = required_state_db(&session)?; let input_path = cwd.join(args.csv_path); let input_path_display = input_path.display().to_string(); @@ -201,7 +206,7 @@ pub async fn handle( })?; if let Err(err) = run_agent_job_loop( session.clone(), - turn.clone(), + environments.to_selections(), db.clone(), job_id.clone(), options, @@ -299,8 +304,10 @@ pub async fn handle( Ok(FunctionToolOutput::from_text(content, Some(true))) } -fn single_local_environment_cwd(turn: &TurnContext) -> Result { - let [turn_environment] = turn.environments.turn_environments.as_slice() else { +fn single_local_environment_cwd( + environments: &TurnEnvironmentSnapshot, +) -> Result { + let [turn_environment] = environments.turn_environments.as_slice() else { return Err(FunctionCallError::RespondToModel( "spawn_agents_on_csv requires exactly one local environment".to_string(), )); diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 9a00fac94dc..8bcbd8ead71 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -9,6 +9,7 @@ use std::time::Instant; use crate::apply_patch; use crate::apply_patch::InternalApplyPatchInvocation; use crate::apply_patch::convert_apply_patch_to_protocol; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -335,6 +336,7 @@ impl ApplyPatchHandler { let ToolInvocation { session, turn, + step, tracker, call_id, tool_name, @@ -360,7 +362,7 @@ impl ApplyPatchHandler { // Verify the parsed patch against the selected environment filesystem. let Some(turn_environment) = - resolve_tool_environment(turn.as_ref(), selected_environment_id.as_deref())? + resolve_tool_environment(step.as_ref(), selected_environment_id.as_deref())? else { return Err(FunctionCallError::RespondToModel( "apply_patch is unavailable in this session".to_string(), @@ -431,6 +433,7 @@ impl ApplyPatchHandler { let tool_ctx = ToolCtx { session: session.clone(), turn: turn.clone(), + environments: step.environments.clone(), call_id: call_id.clone(), tool_name: tool_name.clone(), }; @@ -536,6 +539,7 @@ pub(crate) async fn intercept_apply_patch( turn_environment: TurnEnvironment, session: Arc, turn: Arc, + environments: TurnEnvironmentSnapshot, tracker: Option<&SharedTurnDiffTracker>, call_id: &str, tool_name: &str, @@ -595,6 +599,7 @@ pub(crate) async fn intercept_apply_patch( let tool_ctx = ToolCtx { session: session.clone(), turn: turn.clone(), + environments, call_id: call_id.to_string(), tool_name: ToolName::plain(tool_name), }; diff --git a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs index e9a118641a4..347485fb87f 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch_tests.rs @@ -14,6 +14,7 @@ use tempfile::TempDir; use tokio::sync::Mutex; use crate::session::tests::make_session_and_context; +use crate::session::tests::step_context_for_session; use crate::tools::context::ToolInvocation; use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; @@ -29,9 +30,11 @@ fn sample_patch() -> &'static str { async fn invocation_for_payload(payload: ToolPayload) -> ToolInvocation { let (session, turn) = make_session_and_context().await; + let step = step_context_for_session(&session).await; ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-apply-patch".to_string(), diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index ae79ba405d5..9720cc2f106 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -112,8 +112,8 @@ impl TurnItemEmitter for CoreTurnItemEmitter { async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { let conversation_history = ConversationHistory::new(invocation.session.clone_history().await.into_raw_items()); - let mut environments = Vec::with_capacity(invocation.turn.environments.turn_environments.len()); - for environment in &invocation.turn.environments.turn_environments { + let mut environments = Vec::with_capacity(invocation.step.environments.turn_environments.len()); + for environment in &invocation.step.environments.turn_environments { // TODO(anp): Migrate extension ToolEnvironment and granted-permission lookup to PathUri // so extensions can receive foreign environment cwd values. let Ok(native_cwd) = environment.cwd().to_abs_path() else { @@ -272,9 +272,13 @@ mod tests { async fn exposes_generic_hook_payloads() { let handler = ExtensionToolAdapter::new(Arc::new(StubExtensionExecutor)); let (session, turn) = crate::session::tests::make_session_and_context().await; + let step = Arc::new(crate::session::step_context::StepContext::local_for_test( + &turn, + )); let invocation = ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "call-extension".to_string(), @@ -311,12 +315,13 @@ mod tests { captured_call: Arc::clone(&captured_call), })); let (session, turn, rx) = crate::session::tests::make_session_and_context_with_rx().await; + let step = crate::session::tests::step_context_for_session(session.as_ref()).await; let weak_session = Arc::downgrade(&session); let weak_turn = Arc::downgrade(&turn); let turn_id = turn.sub_id.clone(); let model = turn.model_info.slug.clone(); let truncation_policy = turn.model_info.truncation_policy.into(); - let expected_sandbox_cwds = turn + let expected_sandbox_cwds = step .environments .turn_environments .iter() @@ -342,6 +347,7 @@ mod tests { let invocation = ToolInvocation { session, turn, + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "call-extension".to_string(), @@ -535,6 +541,7 @@ mod tests { async fn image_generation_publication_is_finalized_by_core() { let handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor)); let (session, turn, rx) = crate::session::tests::make_session_and_context_with_rx().await; + let step = crate::session::tests::step_context_for_session(session.as_ref()).await; let expected_path = crate::stream_events_utils::image_generation_artifact_path( &turn.config.codex_home, &session.thread_id.to_string(), @@ -543,6 +550,7 @@ mod tests { let invocation = ToolInvocation { session, turn, + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "call-image".to_string(), diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 016a0344cda..d1d07d7af8d 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -125,6 +125,7 @@ impl McpHandler { let ToolInvocation { session, turn, + step, call_id, payload, .. @@ -143,6 +144,7 @@ impl McpHandler { let result = handle_mcp_tool_call( Arc::clone(&session), &turn, + &step, call_id.clone(), self.tool_info.server_name.clone(), self.tool_info.tool.name.to_string(), @@ -317,6 +319,7 @@ mod search_tests; #[cfg(test)] mod tests { use super::*; + use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::context::ToolCallSource; use crate::tools::hook_names::HookToolName; @@ -340,12 +343,14 @@ mod tests { .to_string(), }; let (session, turn) = make_session_and_context().await; + let step = Arc::new(StepContext::local_for_test(&turn)); let handler = McpHandler::new(tool_info("memory", "memory", "create_entities")) .expect("MCP tool spec should build"); assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-mcp-pre".to_string(), @@ -371,6 +376,7 @@ mod tests { arguments: json!({ "message": "hello" }).to_string(), }; let (session, turn) = make_session_and_context().await; + let step = Arc::new(StepContext::local_for_test(&turn)); let handler = McpHandler::new(tool_info("foo", "mcp__foo", "exec_command")) .expect("MCP tool spec should build"); @@ -378,6 +384,7 @@ mod tests { handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-mcp-pre-builtin-like".to_string(), @@ -398,6 +405,7 @@ mod tests { arguments: json!({ "message": "hello" }).to_string(), }; let (session, turn) = make_session_and_context().await; + let step = Arc::new(StepContext::local_for_test(&turn)); let handler = McpHandler::new(tool_info("foo", "mcp__foo", "exec_command")) .expect("MCP tool spec should build"); @@ -406,6 +414,7 @@ mod tests { ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-mcp-rewrite-builtin-like".to_string(), @@ -448,11 +457,13 @@ mod tests { truncation_policy: codex_utils_output_truncation::TruncationPolicy::Bytes(1024), }; let (session, turn) = make_session_and_context().await; + let step = Arc::new(StepContext::local_for_test(&turn)); let handler = McpHandler::new(tool_info("filesystem", "filesystem", "read_file")) .expect("MCP tool spec should build"); let invocation = ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-mcp-post".to_string(), diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index d9c338e82d7..81bed5de494 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -48,7 +48,7 @@ use std::path::Path; use crate::function_tool::FunctionCallError; use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; -use crate::session::turn_context::TurnContext; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnEnvironment; pub(crate) use crate::tools::code_mode::CodeModeExecuteHandler; pub(crate) use crate::tools::code_mode::CodeModeWaitHandler; @@ -151,13 +151,13 @@ fn resolve_workdir_base_path( } fn resolve_tool_environment<'a>( - turn: &'a TurnContext, + step: &'a StepContext, environment_id: Option<&str>, ) -> Result, FunctionCallError> { environment_id.map_or_else( - || Ok(turn.environments.primary()), + || Ok(step.environments.primary()), |environment_id| { - turn.environments + step.environments .turn_environments .iter() .find(|environment| environment.environment_id == environment_id) diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 204385934e2..eb851e2bcb7 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -49,6 +49,7 @@ async fn handle_spawn_agent( let ToolInvocation { session, turn, + step, payload, call_id, .. @@ -131,7 +132,7 @@ async fn handle_spawn_agent( fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory), parent_thread_id: Some(session.thread_id), - environments: Some(turn.environments.to_selections()), + environments: Some(step.environments.to_selections()), }, )) .await diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 557164f4dea..add45563493 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -74,6 +74,9 @@ fn invocation( ) -> ToolInvocation { ToolInvocation { session, + step: Arc::new(crate::session::step_context::StepContext::local_for_test( + turn.as_ref(), + )), turn, cancellation_token: CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index 51813c5c199..a4a75230103 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -43,6 +43,7 @@ async fn handle_spawn_agent( let ToolInvocation { session, turn, + step, payload, call_id, .. @@ -133,7 +134,7 @@ async fn handle_spawn_agent( fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.clone()), fork_mode, parent_thread_id: Some(session.thread_id), - environments: Some(turn.environments.to_selections()), + environments: Some(step.environments.to_selections()), }, ), ) diff --git a/codex-rs/core/src/tools/handlers/request_permissions.rs b/codex-rs/core/src/tools/handlers/request_permissions.rs index afc15f96597..6952a813c2b 100644 --- a/codex-rs/core/src/tools/handlers/request_permissions.rs +++ b/codex-rs/core/src/tools/handlers/request_permissions.rs @@ -47,6 +47,7 @@ impl RequestPermissionsHandler { let ToolInvocation { session, turn, + step, cancellation_token, call_id, payload, @@ -64,7 +65,7 @@ impl RequestPermissionsHandler { let environment_args: RequestPermissionsEnvironmentArgs = parse_arguments(&arguments)?; let Some(turn_environment) = - resolve_tool_environment(turn.as_ref(), environment_args.environment_id.as_deref())? + resolve_tool_environment(step.as_ref(), environment_args.environment_id.as_deref())? else { return Err(FunctionCallError::RespondToModel( "request_permissions requires a primary environment".to_string(), @@ -92,6 +93,7 @@ impl RequestPermissionsHandler { let response = session .request_permissions_for_environment( &turn, + &step.environments, call_id, args, turn_environment.selection(), diff --git a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs index 7c577c54bef..d743ed0247c 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs @@ -21,6 +21,9 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { agent_nickname: None, agent_role: None, }); + let step = Arc::new(crate::session::step_context::StepContext::local_for_test( + &turn, + )); let result = RequestUserInputHandler { available_modes: Vec::new(), @@ -28,6 +31,7 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { .handle(ToolInvocation { session: Arc::new(session), turn: Arc::new(turn), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), call_id: "call-1".to_string(), diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 3e1f160e9c5..3e67fa5eb70 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -7,6 +7,7 @@ use tokio_util::sync::CancellationToken; use crate::exec::ExecParams; use crate::exec_policy::ExecApprovalRequest; use crate::function_tool::FunctionCallError; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::shell::ShellType; use crate::tools::context::FunctionToolOutput; @@ -52,6 +53,7 @@ struct RunExecLikeArgs { prefix_rule: Option>, session: Arc, turn: Arc, + step: Arc, tracker: crate::tools::context::SharedTurnDiffTracker, call_id: String, shell_runtime_backend: ShellRuntimeBackend, @@ -68,12 +70,13 @@ async fn run_exec_like(args: RunExecLikeArgs) -> Result Result Result Result { let shell = session.user_shell(); let use_login_shell = Self::resolve_use_login_shell(params.login, allow_login_shell)?; let command = Self::base_command(shell.as_ref(), ¶ms.command, use_login_shell); - #[allow(deprecated)] - let cwd = turn_context.resolve_path(params.workdir.clone()); - Ok(ExecParams { command, cwd, @@ -156,6 +155,7 @@ impl ShellCommandHandler { let ToolInvocation { session, turn, + step, cancellation_token, tracker, call_id, @@ -170,16 +170,13 @@ impl ShellCommandHandler { ))); }; - #[allow(deprecated)] - let cwd = resolve_workdir_base_path(&arguments, &turn.cwd)?; + let cwd = resolve_workdir_base_path(&arguments, &step.effective_cwd(turn.as_ref()))?; let params: ShellCommandToolCallParams = parse_arguments_with_base_path(&arguments, &cwd)?; - #[allow(deprecated)] - let workdir = turn.resolve_path(params.workdir.clone()); maybe_emit_implicit_skill_invocation( session.as_ref(), turn.as_ref(), ¶ms.command, - &workdir, + &cwd, ) .await; let prefix_rule = params.prefix_rule.clone(); @@ -187,6 +184,7 @@ impl ShellCommandHandler { ¶ms, session.as_ref(), turn.as_ref(), + cwd, session.thread_id, turn.config.permissions.allow_login_shell, )?; @@ -201,6 +199,7 @@ impl ShellCommandHandler { prefix_rule, session, turn, + step, tracker, call_id, shell_runtime_backend: self.shell_runtime_backend(), diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index 2ec5cdf3ede..b12616071eb 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -81,7 +81,10 @@ async fn shell_command_handler_to_exec_params_uses_session_shell_and_turn_contex .user_shell() .derive_exec_args(&command, /*use_login_shell*/ true); #[allow(deprecated)] - let expected_cwd = turn_context.resolve_path(workdir.clone()); + let expected_cwd = workdir.as_ref().map_or_else( + || turn_context.cwd.clone(), + |path| turn_context.cwd.join(path), + ); let expected_env = create_env( &turn_context.config.permissions.shell_environment_policy, Some(session.thread_id), @@ -102,6 +105,7 @@ async fn shell_command_handler_to_exec_params_uses_session_shell_and_turn_contex ¶ms, &session, &turn_context, + expected_cwd.clone(), session.thread_id, /*allow_login_shell*/ true, ) @@ -164,6 +168,8 @@ async fn shell_command_handler_defaults_to_non_login_when_disallowed() { ¶ms, &session, &turn_context, + #[allow(deprecated)] + turn_context.cwd.clone(), session.thread_id, /*allow_login_shell*/ false, ) @@ -196,12 +202,16 @@ async fn shell_command_pre_tool_use_payload_uses_raw_command() { arguments: json!({ "command": "printf shell command" }).to_string(), }; let (session, turn) = make_session_and_context().await; + let step = Arc::new(crate::session::step_context::StepContext::local_for_test( + &turn, + )); let handler = ShellCommandHandler::from(codex_tools::ShellCommandBackendConfig::Classic); assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-42".to_string(), @@ -228,9 +238,13 @@ async fn build_post_tool_use_payload_uses_tool_output_wire_value() { }; let handler = ShellCommandHandler::from(codex_tools::ShellCommandBackendConfig::Classic); let (session, turn) = make_session_and_context().await; + let step = Arc::new(crate::session::step_context::StepContext::local_for_test( + &turn, + )); let invocation = ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-42".to_string(), diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index db13c6f0f9f..a79bd328804 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -108,6 +108,7 @@ impl ExecCommandHandler { let ToolInvocation { session, turn, + step, tracker, call_id, payload, @@ -124,10 +125,15 @@ impl ExecCommandHandler { }; let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager; - let context = UnifiedExecContext::new(session.clone(), turn.clone(), call_id.clone()); + let context = UnifiedExecContext::new( + session.clone(), + turn.clone(), + step.environments.clone(), + call_id.clone(), + ); let environment_args: ExecCommandEnvironmentArgs = parse_arguments(&arguments)?; let Some(turn_environment) = - resolve_tool_environment(turn.as_ref(), environment_args.environment_id.as_deref())? + resolve_tool_environment(step.as_ref(), environment_args.environment_id.as_deref())? else { return Err(FunctionCallError::RespondToModel( "unified exec is unavailable in this session".to_string(), @@ -296,6 +302,7 @@ impl ExecCommandHandler { turn_environment.clone(), context.session.clone(), context.turn.clone(), + context.environments.clone(), Some(&tracker), &context.call_id, "exec_command", diff --git a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs index b33a0e85b0f..0663aeb939b 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -10,6 +10,7 @@ use pretty_assertions::assert_eq; use std::sync::Arc; use crate::session::tests::make_session_and_context; +use crate::session::tests::step_context_for_session; use crate::tools::context::ExecCommandToolOutput; use crate::tools::context::ToolCallSource; use crate::tools::context::ToolInvocation; @@ -27,9 +28,11 @@ async fn invocation_for_payload( payload: ToolPayload, ) -> ToolInvocation { let (session, turn) = make_session_and_context().await; + let step = step_context_for_session(&session).await; ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: call_id.to_string(), @@ -236,12 +239,16 @@ async fn exec_command_pre_tool_use_payload_uses_raw_command() { arguments: serde_json::json!({ "cmd": "printf exec command" }).to_string(), }; let (session, turn) = make_session_and_context().await; + let step = Arc::new(crate::session::step_context::StepContext::local_for_test( + &turn, + )); let handler = ExecCommandHandler::default(); assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-43".to_string(), @@ -262,12 +269,16 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() { arguments: serde_json::json!({ "chars": "echo hi" }).to_string(), }; let (session, turn) = make_session_and_context().await; + let step = Arc::new(crate::session::step_context::StepContext::local_for_test( + &turn, + )); let handler = WriteStdinHandler; assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { session: session.into(), turn: turn.into(), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-44".to_string(), diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 84911fb8150..182d5939119 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -104,6 +104,7 @@ impl ViewImageHandler { let ToolInvocation { session, turn, + step, payload, call_id, .. @@ -137,7 +138,7 @@ impl ViewImageHandler { }; let Some(turn_environment) = - resolve_tool_environment(turn.as_ref(), environment_id.as_deref())? + resolve_tool_environment(step.as_ref(), environment_id.as_deref())? else { return Err(FunctionCallError::RespondToModel( "view_image is unavailable in this session".to_string(), @@ -272,6 +273,7 @@ impl ToolOutput for ViewImageOutput { #[cfg(test)] mod tests { use super::*; + use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnEnvironment; use crate::tools::context::ToolCallSource; @@ -286,14 +288,14 @@ mod tests { use std::sync::Arc; use tokio::sync::Mutex; - fn replace_primary_environment_cwd(turn: &mut crate::TurnContext, cwd: AbsolutePathBuf) { - let current = turn + fn replace_primary_environment_cwd(step: &mut StepContext, cwd: AbsolutePathBuf) { + let current = step .environments .turn_environments .first() .cloned() .expect("default local turn environment"); - turn.environments.turn_environments[0] = TurnEnvironment::new( + step.environments.turn_environments[0] = TurnEnvironment::new( current.environment_id, current.environment, PathUri::from_abs_path(&cwd), @@ -337,7 +339,8 @@ mod tests { let image_dir = tempfile::tempdir().expect("create image temp dir"); let image_cwd = image_dir.abs(); - replace_primary_environment_cwd(&mut turn, image_cwd.clone()); + let mut step = StepContext::local_for_test(&turn); + replace_primary_environment_cwd(&mut step, image_cwd.clone()); let image_path = image_cwd.join("image.png"); std::fs::write(image_path.as_path(), b"not a real image").expect("write test image"); turn.permission_profile = PermissionProfile::read_only(); @@ -346,6 +349,7 @@ mod tests { .handle(ToolInvocation { session: Arc::new(session), turn: Arc::new(turn), + step: Arc::new(step), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-view-image".to_string(), @@ -369,11 +373,13 @@ mod tests { #[tokio::test] async fn handle_rejects_unsupported_detail() { let (session, turn) = make_session_and_context().await; + let step = Arc::new(StepContext::local_for_test(&turn)); let result = ViewImageHandler::default() .handle(ToolInvocation { session: Arc::new(session), turn: Arc::new(turn), + step, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-view-image".to_string(), @@ -400,7 +406,8 @@ mod tests { let image_dir = tempfile::tempdir().expect("create image temp dir"); let image_cwd = image_dir.abs(); - replace_primary_environment_cwd(&mut turn, image_cwd.clone()); + let mut step = StepContext::local_for_test(&turn); + replace_primary_environment_cwd(&mut step, image_cwd.clone()); let image_path = image_cwd.join("image.png"); std::fs::write(image_path.as_path(), b"not a real image").expect("write test image"); turn.permission_profile = PermissionProfile::Disabled; @@ -409,6 +416,7 @@ mod tests { .handle(ToolInvocation { session: Arc::new(session), turn: Arc::new(turn), + step: Arc::new(step), cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-view-image".to_string(), diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index b05a7455267..581f6c017a9 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -499,9 +499,13 @@ impl NetworkApprovalService { let use_guardian = routes_approval_to_guardian(&turn_context); let guardian_review_id = use_guardian.then(new_guardian_review_id); let approval_decision = if let Some(review_id) = guardian_review_id.clone() { + // V0 uses the latest environment for asynchronous network review. Preserve the + // originating environment once blocked proxy requests identify their tool call. + let environments = session.services.turn_environments.snapshot().await; review_approval_request( &session, &turn_context, + environments, review_id, GuardianApprovalRequest::NetworkAccess { id: guardian_approval_id.clone(), diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index 803b8d6a9fd..87d46838a47 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -79,6 +79,7 @@ impl ToolOrchestrator { let attempt_tool_ctx = ToolCtx { session: tool_ctx.session.clone(), turn: tool_ctx.turn.clone(), + environments: tool_ctx.environments.clone(), call_id: tool_ctx.call_id.clone(), tool_name: tool_ctx.tool_name.clone(), }; @@ -161,6 +162,7 @@ impl ToolOrchestrator { let approval_ctx = ApprovalCtx { session: &tool_ctx.session, turn: &tool_ctx.turn, + environments: &tool_ctx.environments, call_id: &tool_ctx.call_id, guardian_review_id: guardian_review_id.clone(), retry_reason: None, @@ -196,6 +198,7 @@ impl ToolOrchestrator { let approval_ctx = ApprovalCtx { session: &tool_ctx.session, turn: &tool_ctx.turn, + environments: &tool_ctx.environments, call_id: &tool_ctx.call_id, guardian_review_id: guardian_review_id.clone(), retry_reason: reason.clone(), @@ -379,6 +382,7 @@ impl ToolOrchestrator { let approval_ctx = ApprovalCtx { session: &tool_ctx.session, turn: &tool_ctx.turn, + environments: &tool_ctx.environments, call_id: &tool_ctx.call_id, guardian_review_id: guardian_review_id.clone(), retry_reason: Some(retry_reason), diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index c40886b56cb..7203755c685 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -240,6 +240,7 @@ mod tests { use super::*; use std::time::Duration; + use crate::session::step_context::StepContext; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::registry::CoreToolRuntime; @@ -423,9 +424,11 @@ mod tests { let handler = Arc::new(ImmediateHandler { tool_name: tool_name.clone(), }) as Arc; + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let router = Arc::new(ToolRouter::from_parts( ToolRegistry::from_tools([handler]), Vec::new(), + step_context, )); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let runtime = ToolCallRuntime::new(router, session, turn_context, tracker); @@ -495,9 +498,11 @@ mod tests { cleanup_started: std::sync::Mutex::new(Some(cleanup_started_tx)), allow_cleanup: Arc::clone(&allow_cleanup), }) as Arc; + let step_context = Arc::new(StepContext::local_for_test(turn_context.as_ref())); let router = Arc::new(ToolRouter::from_parts( ToolRegistry::from_tools([handler]), Vec::new(), + step_context, )); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let runtime = ToolCallRuntime::new(router, session, turn_context, tracker); diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index 67613f1d081..3e5dbae5da2 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -459,6 +459,9 @@ fn test_invocation( ) -> ToolInvocation { ToolInvocation { session, + step: Arc::new(crate::session::step_context::StepContext::local_for_test( + turn.as_ref(), + )), turn, cancellation_token: tokio_util::sync::CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new( diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index b4884483984..0f916b42f4f 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -1,5 +1,6 @@ use crate::function_tool::FunctionCallError; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolInvocation; @@ -35,6 +36,7 @@ pub struct ToolCall { pub struct ToolRouter { registry: ToolRegistry, model_visible_specs: Vec, + step_context: Arc, } pub(crate) struct ToolRouterParams<'a> { @@ -58,18 +60,29 @@ pub(crate) struct ToolSuggestCandidates { } impl ToolRouter { - pub(crate) fn from_turn_context( + pub(crate) fn from_contexts( turn_context: &TurnContext, + step_context: Arc, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> Self { - build_tool_router(turn_context, params, tool_search_handler_cache) + build_tool_router( + turn_context, + step_context, + params, + tool_search_handler_cache, + ) } - pub(crate) fn from_parts(registry: ToolRegistry, model_visible_specs: Vec) -> Self { + pub(crate) fn from_parts( + registry: ToolRegistry, + model_visible_specs: Vec, + step_context: Arc, + ) -> Self { Self { registry, model_visible_specs, + step_context, } } @@ -160,6 +173,7 @@ impl ToolRouter { } #[allow(dead_code)] + #[allow(clippy::too_many_arguments)] #[instrument(level = "trace", skip_all, err)] pub async fn dispatch_tool_call_with_code_mode_result( &self, @@ -226,6 +240,7 @@ impl ToolRouter { let invocation = ToolInvocation { session, turn, + step: Arc::clone(&self.step_context), cancellation_token, tracker, call_id, diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index a38ab8af26c..a4116be63e1 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use crate::config::Config; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::tools::context::ToolPayload; use crate::turn_diff_tracker::TurnDiffTracker; @@ -111,8 +112,10 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow .load_full() .list_all_tools() .await; - let router = ToolRouter::from_turn_context( + let step = Arc::new(StepContext::local_for_test(&turn)); + let router = ToolRouter::from_contexts( &turn, + step, ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -179,8 +182,10 @@ async fn build_tool_call_uses_namespace_for_registry_name() -> anyhow::Result<() #[tokio::test] async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; - let router = ToolRouter::from_turn_context( + let step = Arc::new(StepContext::local_for_test(&turn)); + let router = ToolRouter::from_contexts( &turn, + step, ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -228,8 +233,10 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { #[tokio::test] async fn tools_without_handlers_do_not_support_parallel() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; - let router = ToolRouter::from_turn_context( + let step = Arc::new(StepContext::local_for_test(&turn)); + let router = ToolRouter::from_contexts( &turn, + step, ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -283,8 +290,10 @@ async fn specs_filter_deferred_dynamic_tools() -> anyhow::Result<()> { ], })]; - let router = ToolRouter::from_turn_context( + let step = Arc::new(StepContext::local_for_test(&turn)); + let router = ToolRouter::from_contexts( &turn, + step, ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -346,8 +355,10 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow .record_conversation_items(&turn, std::slice::from_ref(&history_item)) .await; - let router = ToolRouter::from_turn_context( + let step = Arc::new(StepContext::local_for_test(&turn)); + let router = ToolRouter::from_contexts( &turn, + step, ToolRouterParams { tool_suggest_candidates: None, deferred_mcp_tools: None, @@ -379,10 +390,11 @@ async fn extension_tool_executors_are_model_visible_and_dispatchable() -> anyhow metadata: None, })? .expect("function_call should produce a tool call"); + let turn = Arc::new(turn); let result = router .dispatch_tool_call_with_code_mode_result( Arc::new(session), - Arc::new(turn), + turn, CancellationToken::new(), Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call, diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index 9611429a7dd..76f4fe7475b 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -161,8 +161,15 @@ impl Approvable for ApplyPatchRuntime { return ReviewDecision::Abort; } }; - return review_approval_request(session, turn, review_id, action, retry_reason) - .await; + return review_approval_request( + session, + turn, + ctx.environments.clone(), + review_id, + action, + retry_reason, + ) + .await; } if req.permissions_preapproved && retry_reason.is_none() { return ReviewDecision::Approved; diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index f88c2be3d4d..4baba6cd118 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -156,6 +156,7 @@ impl Approvable for ShellRuntime { return review_approval_request( session, turn, + ctx.environments.clone(), review_id, GuardianApprovalRequest::Shell { id: call_id, diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index b8e417f8565..55949ee893b 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -229,6 +229,7 @@ pub(super) async fn try_run_zsh_fork( policy: Arc::clone(&exec_policy), session: Arc::clone(&ctx.session), turn: Arc::clone(&ctx.turn), + environments: ctx.environments.clone(), call_id: ctx.call_id.clone(), environment_id: req.turn_environment.environment_id.clone(), tool_name: GuardianCommandSource::Shell, @@ -312,6 +313,7 @@ pub(crate) async fn prepare_unified_exec_zsh_fork( policy: Arc::clone(&exec_policy), session: Arc::clone(&ctx.session), turn: Arc::clone(&ctx.turn), + environments: ctx.environments.clone(), call_id: ctx.call_id.clone(), environment_id: req.turn_environment.environment_id.clone(), tool_name: GuardianCommandSource::UnifiedExec, @@ -347,6 +349,7 @@ struct CoreShellActionProvider { policy: Arc>, session: Arc, turn: Arc, + environments: crate::environment_selection::TurnEnvironmentSnapshot, call_id: String, environment_id: String, tool_name: GuardianCommandSource, @@ -443,6 +446,7 @@ impl CoreShellActionProvider { let workdir = workdir.clone(); let session = self.session.clone(); let turn = self.turn.clone(); + let environments = self.environments.clone(); let call_id = self.call_id.clone(); let approval_id = Some(Uuid::new_v4().to_string()); let environment_id = Some(self.environment_id.clone()); @@ -486,6 +490,7 @@ impl CoreShellActionProvider { let decision = review_approval_request( &session, &turn, + environments, review_id.clone(), GuardianApprovalRequest::Execve { id: call_id.clone(), diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs index 2b001476129..405c464a73e 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs @@ -423,6 +423,8 @@ async fn preapproved_additional_permissions_escalate_intercepted_exec() -> anyho let provider = CoreShellActionProvider { policy: Arc::new(RwLock::new(codex_execpolicy::Policy::empty())), session: Arc::new(session), + environments: crate::session::step_context::StepContext::local_for_test(&turn_context) + .environments, turn: Arc::new(turn_context), call_id: "preapproved-additional-permissions".to_string(), environment_id: "local".to_string(), @@ -559,6 +561,8 @@ async fn execve_permission_request_hook_short_circuits_prompt() -> anyhow::Resul let provider = CoreShellActionProvider { policy: std::sync::Arc::new(RwLock::new(codex_execpolicy::Policy::empty())), session: std::sync::Arc::new(session), + environments: crate::session::step_context::StepContext::local_for_test(&turn_context) + .environments, turn: std::sync::Arc::new(turn_context), call_id: "execve-hook-call".to_string(), environment_id: "local".to_string(), @@ -770,6 +774,8 @@ prefix_rule(pattern = ["{cat_path_literal}"], decision = "allow") let provider = CoreShellActionProvider { policy: Arc::new(RwLock::new(policy)), session: Arc::new(session), + environments: crate::session::step_context::StepContext::local_for_test(&turn_context) + .environments, turn: Arc::new(turn_context), call_id: "deny-read-prefix-allow".to_string(), environment_id: "local".to_string(), @@ -807,6 +813,8 @@ async fn denied_reads_keep_granular_sandbox_rejection_for_escalation() -> anyhow let provider = CoreShellActionProvider { policy: Arc::new(RwLock::new(PolicyParser::new().build())), session: Arc::new(session), + environments: crate::session::step_context::StepContext::local_for_test(&turn_context) + .environments, turn: Arc::new(turn_context), call_id: "deny-read-granular-sandbox-reject".to_string(), environment_id: "local".to_string(), diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 80f0990b5c9..6636e34c2eb 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -192,6 +192,7 @@ impl Approvable for UnifiedExecRuntime<'_> { return review_approval_request( session, turn, + ctx.environments.clone(), review_id, GuardianApprovalRequest::ExecCommand { id: call_id, diff --git a/codex-rs/core/src/tools/sandboxing.rs b/codex-rs/core/src/tools/sandboxing.rs index cc7a7913ddc..6bba5f0830d 100644 --- a/codex-rs/core/src/tools/sandboxing.rs +++ b/codex-rs/core/src/tools/sandboxing.rs @@ -4,6 +4,7 @@ //! `ApprovalCtx`, `Approvable`) together with the sandbox orchestration traits //! and helpers (`Sandboxable`, `ToolRuntime`, `SandboxAttempt`, etc.). +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::sandboxing::ExecOptions; use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; @@ -119,6 +120,7 @@ where pub(crate) struct ApprovalCtx<'a> { pub session: &'a Arc, pub turn: &'a Arc, + pub environments: &'a TurnEnvironmentSnapshot, pub call_id: &'a str, /// Guardian review lifecycle ID for this approval, when guardian is reviewing it. /// @@ -379,6 +381,7 @@ pub(crate) trait Sandboxable { pub(crate) struct ToolCtx { pub session: Arc, pub turn: Arc, + pub environments: TurnEnvironmentSnapshot, pub call_id: String, pub tool_name: ToolName, } diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index c697f287aa6..9739bb00441 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -1,5 +1,6 @@ use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::code_mode::execute_spec::create_code_mode_tool; use crate::tools::context::ToolInvocation; @@ -143,6 +144,7 @@ impl PlannedTools { #[derive(Clone, Copy)] struct CoreToolPlanContext<'a> { turn_context: &'a TurnContext, + step_context: &'a StepContext, mcp_tools: Option<&'a [ToolInfo]>, deferred_mcp_tools: Option<&'a [ToolInfo]>, tool_suggest_candidates: Option<&'a crate::tools::router::ToolSuggestCandidates>, @@ -156,17 +158,23 @@ struct CoreToolPlanContext<'a> { #[instrument(level = "trace", skip_all)] pub(crate) fn build_tool_router( turn_context: &TurnContext, + step_context: Arc, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> ToolRouter { - let (model_visible_specs, registry) = - build_tool_specs_and_registry(turn_context, params, tool_search_handler_cache); - ToolRouter::from_parts(registry, model_visible_specs) + let (model_visible_specs, registry) = build_tool_specs_and_registry( + turn_context, + step_context.as_ref(), + params, + tool_search_handler_cache, + ); + ToolRouter::from_parts(registry, model_visible_specs, step_context) } #[instrument(level = "trace", skip_all)] fn build_tool_specs_and_registry( turn_context: &TurnContext, + step_context: &StepContext, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> (Vec, ToolRegistry) { @@ -181,6 +189,7 @@ fn build_tool_specs_and_registry( crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let context = CoreToolPlanContext { turn_context, + step_context, mcp_tools: mcp_tools.as_deref(), deferred_mcp_tools: deferred_mcp_tools.as_deref(), tool_suggest_candidates: tool_suggest_candidates.as_ref(), @@ -598,7 +607,7 @@ fn standalone_web_search_enabled(turn_context: &TurnContext) -> bool { fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { let turn_context = context.turn_context; let features = turn_context.config.features.get(); - let environment_mode = turn_context.tool_environment_mode(); + let environment_mode = context.step_context.tool_environment_mode(); if !environment_mode.has_environment() { return; } @@ -618,7 +627,7 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Planne allow_login_shell, exec_permission_approvals_enabled, include_environment_id, - include_shell_parameter: unified_exec_should_include_shell_parameter(turn_context), + include_shell_parameter: unified_exec_should_include_shell_parameter(context), })); planned_tools.add(WriteStdinHandler); @@ -635,11 +644,12 @@ fn add_shell_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Planne } } -fn unified_exec_should_include_shell_parameter(turn_context: &TurnContext) -> bool { +fn unified_exec_should_include_shell_parameter(context: &CoreToolPlanContext<'_>) -> bool { !matches!( - &turn_context.unified_exec_shell_mode, + &context.turn_context.unified_exec_shell_mode, UnifiedExecShellMode::ZshFork(_) - ) || turn_context + ) || context + .step_context .environments .turn_environments .iter() @@ -659,7 +669,7 @@ fn add_mcp_resource_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { let turn_context = context.turn_context; let features = turn_context.config.features.get(); - let environment_mode = turn_context.tool_environment_mode(); + let environment_mode = context.step_context.tool_environment_mode(); planned_tools.add(PlanHandler); diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 43bcf03a6cc..f336d2c93b0 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -30,6 +30,7 @@ use codex_tools::ToolSpec; use pretty_assertions::assert_eq; use serde_json::json; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; use crate::tools::handlers::ToolSearchHandlerCache; @@ -176,10 +177,21 @@ async fn probe_with( configure_turn: impl FnOnce(&mut TurnContext), inputs: ToolPlanInputs, ) -> ToolPlanProbe { - let (_session, mut turn) = make_session_and_context().await; + probe_with_step(configure_turn, |_| {}, inputs).await +} + +async fn probe_with_step( + configure_turn: impl FnOnce(&mut TurnContext), + configure_step: impl FnOnce(&mut StepContext), + inputs: ToolPlanInputs, +) -> ToolPlanProbe { + let (session, mut turn) = make_session_and_context().await; + let mut step = StepContext::new(session.services.turn_environments.snapshot().await); configure_turn(&mut turn); - let router = ToolRouter::from_turn_context( + configure_step(&mut step); + let router = ToolRouter::from_contexts( &turn, + Arc::new(step), ToolRouterParams { tool_suggest_candidates: inputs.tool_suggest_candidates, mcp_tools: inputs.mcp_tools, @@ -331,10 +343,10 @@ impl ToolExecutor for DeferredExtensionTool { } } -fn duplicate_primary_environment(turn: &mut TurnContext) { - let mut second_environment = turn.environments.turn_environments[0].clone(); +fn duplicate_primary_environment(step: &mut StepContext) { + let mut second_environment = step.environments.turn_environments[0].clone(); second_environment.environment_id = "secondary".to_string(); - turn.environments.turn_environments.push(second_environment); + step.environments.turn_environments.push(second_environment); } fn mcp_tool(server: &str, namespace: &str, name: &str) -> ToolInfo { @@ -555,38 +567,43 @@ async fn zsh_fork_unified_exec_keeps_shell_parameter_when_remote_environment_ava return; } - let plan = probe(|turn| { - set_features( - turn, - &[ - Feature::ShellTool, - Feature::UnifiedExec, - Feature::ShellZshFork, - Feature::UnifiedExecZshFork, - ], - ); - turn.unified_exec_shell_mode = - codex_tools::UnifiedExecShellMode::ZshFork(zsh_fork_config_for_spec_plan_tests()); - let remote_cwd = turn - .environments - .primary() - .expect("primary environment") - .cwd() - .clone(); - turn.environments.turn_environments.push( - crate::session::turn_context::TurnEnvironment::new( - "remote".to_string(), - Arc::new( - codex_exec_server::Environment::create_for_tests(Some( - "ws://127.0.0.1:1/remote-exec-server".to_string(), - )) - .expect("remote test environment"), + let plan = probe_with_step( + |turn| { + set_features( + turn, + &[ + Feature::ShellTool, + Feature::UnifiedExec, + Feature::ShellZshFork, + Feature::UnifiedExecZshFork, + ], + ); + turn.unified_exec_shell_mode = + codex_tools::UnifiedExecShellMode::ZshFork(zsh_fork_config_for_spec_plan_tests()); + }, + |step| { + let remote_cwd = step + .environments + .primary() + .expect("primary environment") + .cwd() + .clone(); + step.environments.turn_environments.push( + crate::session::turn_context::TurnEnvironment::new( + "remote".to_string(), + Arc::new( + codex_exec_server::Environment::create_for_tests(Some( + "ws://127.0.0.1:1/remote-exec-server".to_string(), + )) + .expect("remote test environment"), + ), + remote_cwd, + /*shell*/ None, ), - remote_cwd, - /*shell*/ None, - ), - ); - }) + ); + }, + ToolPlanInputs::default(), + ) .await; plan.assert_visible_contains(&["exec_command", "write_stdin"]); @@ -599,11 +616,14 @@ async fn zsh_fork_unified_exec_keeps_shell_parameter_when_remote_environment_ava #[tokio::test] async fn environment_count_controls_environment_backed_tools() { - let no_environment = probe(|turn| { - turn.environments.turn_environments.clear(); - set_feature(turn, Feature::ShellTool, /*enabled*/ true); - turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); - }) + let no_environment = probe_with_step( + |turn| { + set_feature(turn, Feature::ShellTool, /*enabled*/ true); + turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); + }, + |step| step.environments.turn_environments.clear(), + ToolPlanInputs::default(), + ) .await; no_environment.assert_visible_lacks(&[ "shell_command", @@ -618,12 +638,15 @@ async fn environment_count_controls_environment_backed_tools() { "view_image", ]); - let multiple_environments = probe(|turn| { - duplicate_primary_environment(turn); - set_feature(turn, Feature::ShellTool, /*enabled*/ true); - set_feature(turn, Feature::UnifiedExec, /*enabled*/ true); - turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); - }) + let multiple_environments = probe_with_step( + |turn| { + set_feature(turn, Feature::ShellTool, /*enabled*/ true); + set_feature(turn, Feature::UnifiedExec, /*enabled*/ true); + turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); + }, + duplicate_primary_environment, + ToolPlanInputs::default(), + ) .await; multiple_environments.assert_visible_contains(&["exec_command", "apply_patch", "view_image"]); assert!(has_parameter( @@ -771,10 +794,14 @@ async fn deferred_extension_tools_are_discoverable_with_tool_search() { async fn tool_search_cache_rebuilds_when_deferred_sources_change() { let cache = ToolSearchHandlerCache::default(); - let (_session, mut first_turn) = make_session_and_context().await; + let (first_session, mut first_turn) = make_session_and_context().await; first_turn.model_info.supports_search_tool = true; - let first_router = ToolRouter::from_turn_context( + let first_step = Arc::new(StepContext::new( + first_session.services.turn_environments.snapshot().await, + )); + let first_router = ToolRouter::from_contexts( &first_turn, + first_step, ToolRouterParams { mcp_tools: None, deferred_mcp_tools: Some(vec![mcp_tool("first", "mcp__first", "lookup")]), @@ -786,10 +813,14 @@ async fn tool_search_cache_rebuilds_when_deferred_sources_change() { ); let first_plan = ToolPlanProbe::from_router(first_router); - let (_session, mut second_turn) = make_session_and_context().await; + let (second_session, mut second_turn) = make_session_and_context().await; second_turn.model_info.supports_search_tool = true; - let second_router = ToolRouter::from_turn_context( + let second_step = Arc::new(StepContext::new( + second_session.services.turn_environments.snapshot().await, + )); + let second_router = ToolRouter::from_contexts( &second_turn, + second_step, ToolRouterParams { mcp_tools: None, deferred_mcp_tools: Some(vec![mcp_tool("second", "mcp__second", "lookup")]), diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index b09996872ea..830f5dd1bf4 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -271,6 +271,9 @@ fn test_invocation_with_payload( ) -> ToolInvocation { ToolInvocation { session, + step: Arc::new(crate::session::step_context::StepContext::local_for_test( + turn.as_ref(), + )), turn, cancellation_token: CancellationToken::new(), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 34d88c81491..65a2d0c7d60 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -36,6 +36,7 @@ use rand::Rng; use rand::rng; use tokio::sync::Mutex; +use crate::environment_selection::TurnEnvironmentSnapshot; use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -75,14 +76,21 @@ pub(crate) const MAX_UNIFIED_EXEC_PROCESSES: usize = 64; pub(crate) struct UnifiedExecContext { pub session: Arc, pub turn: Arc, + pub environments: TurnEnvironmentSnapshot, pub call_id: String, } impl UnifiedExecContext { - pub fn new(session: Arc, turn: Arc, call_id: String) -> Self { + pub fn new( + session: Arc, + turn: Arc, + environments: TurnEnvironmentSnapshot, + call_id: String, + ) -> Self { Self { session, turn, + environments, call_id, } } diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 0d4e21d93c6..e8dca443893 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -5,6 +5,7 @@ use crate::exec::ExecCapturePolicy; use crate::exec::ExecExpiration; use crate::sandboxing::ExecRequest; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; use crate::tools::context::ExecCommandToolOutput; @@ -96,6 +97,7 @@ async fn exec_command_with_tty( workdir: Option, tty: bool, ) -> Result { + let step = Arc::new(StepContext::local_for_test(turn.as_ref())); let manager = &session.services.unified_exec_manager; let process_id = manager.allocate_process_id().await; #[allow(deprecated)] @@ -112,7 +114,7 @@ async fn exec_command_with_tty( &request, tty, Box::new(NoopSpawnLifecycle), - turn.environments + step.environments .primary() .expect("turn environment") .environment @@ -120,8 +122,12 @@ async fn exec_command_with_tty( ) .await?, ); - let context = - UnifiedExecContext::new(Arc::clone(session), Arc::clone(turn), "call".to_string()); + let context = UnifiedExecContext::new( + Arc::clone(session), + Arc::clone(turn), + step.environments.clone(), + "call".to_string(), + ); let started_at = Instant::now(); let process_started_alive = !process.has_exited() && process.exit_code().is_none(); if process_started_alive { @@ -894,8 +900,9 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<() }; let remote_test_env = remote_test_env().await?; - let (_, mut turn) = make_session_and_context().await; - turn.environments.turn_environments[0].environment = + let (_, turn) = make_session_and_context().await; + let mut step = StepContext::local_for_test(&turn); + step.environments.turn_environments[0].environment = Arc::new(remote_test_env.environment().clone()); #[allow(deprecated)] @@ -916,7 +923,7 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<() Box::new(TestSpawnLifecycle { inherited_fds: vec![42], }), - turn.environments + step.environments .primary() .expect("turn environment") .environment diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 0aac4918232..484f7640fc9 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -1108,6 +1108,7 @@ impl UnifiedExecProcessManager { let tool_ctx = ToolCtx { session: context.session.clone(), turn: context.turn.clone(), + environments: context.environments.clone(), call_id: context.call_id.clone(), tool_name: ToolName::plain("exec_command"), }; diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index c399b9c9887..c3f36fa2ee9 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -183,9 +183,11 @@ async fn late_network_denial_grace_observes_cancellation_after_exit() { #[tokio::test] async fn failed_initial_end_for_unstored_process_uses_fallback_output() { let (session, turn, rx_event) = crate::session::tests::make_session_and_context_with_rx().await; + let step = crate::session::tests::step_context_for_session(session.as_ref()).await; let context = UnifiedExecContext::new( Arc::clone(&session), Arc::clone(&turn), + step.environments.clone(), "call-unified-denied".to_string(), ); let request = ExecCommandRequest { @@ -203,7 +205,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { cwd: turn.cwd.clone().into(), #[allow(deprecated)] sandbox_cwd: turn.cwd.clone().into(), - turn_environment: turn + turn_environment: step .environments .primary() .cloned() diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 60b476d6143..cab1e34efa4 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -23,6 +23,7 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::user_input::UserInput; use codex_web_search_extension::install as install_web_search_extension; +use core_test_support::PathExt; use core_test_support::apps_test_server::AppsTestServer; use core_test_support::apps_test_server::AppsTestToolLoading; use core_test_support::apps_test_server::DIRECT_CALENDAR_APP_ONLY_TOOL; @@ -41,6 +42,7 @@ use core_test_support::skip_if_no_network; use core_test_support::skip_if_wine_exec; use core_test_support::stdio_server_bin; use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::local; use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; @@ -2259,15 +2261,24 @@ async fn code_mode_background_keeps_running_on_later_turn_without_wait() -> Resu let _ = config.features.enable(Feature::CodeMode); }); let test = builder.build(&server).await?; - let resumed_file = test.workspace_path("code-mode-yield-resumed.txt"); + let resumed_file_name = "code-mode-yield-resumed.txt"; + let resumed_file = test.workspace_path(resumed_file_name); + let later_workspace = test.workspace_path("later-workspace"); + fs::create_dir_all(&later_workspace)?; + let later_resumed_file = later_workspace.join(resumed_file_name); + let later_turn_gate = test.workspace_path("later-turn-started.ready"); + let wait_for_later_turn = wait_for_file_source(&later_turn_gate)?; + let later_turn_gate_quoted = shlex::try_join([later_turn_gate.to_string_lossy().as_ref()])?; let resumed_file_quoted = shlex::try_join([resumed_file.to_string_lossy().as_ref()])?; - let write_file_command = format!("printf resumed > {resumed_file_quoted}"); - let wait_for_file_command = - format!("while [ ! -f {resumed_file_quoted} ]; do sleep 0.01; done; printf ready"); + let write_file_command = format!("printf resumed > {resumed_file_name}"); + let wait_for_file_command = format!( + "i=0; while [ \"$i\" -lt 100 ]; do [ -f {resumed_file_quoted} ] && {{ printf ready; exit; }}; i=$((i+1)); sleep 0.01; done; printf missing" + ); let code = format!( r#" text("before yield"); yield_control(); +{wait_for_later_turn} await tools.exec_command({{ cmd: {write_file_command:?} }}); text("after yield"); "# @@ -2313,31 +2324,51 @@ text("after yield"); "call-2", "exec_command", &serde_json::to_string(&serde_json::json!({ - "cmd": wait_for_file_command, + "cmd": format!("touch {later_turn_gate_quoted}"), }))?, ), ev_completed("resp-3"), ]), ) .await; + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-4"), + responses::ev_function_call( + "call-3", + "exec_command", + &serde_json::to_string(&serde_json::json!({ + "cmd": wait_for_file_command, + }))?, + ), + ev_completed("resp-4"), + ]), + ) + .await; let second_completion = responses::mount_sse_once( &server, sse(vec![ ev_assistant_message("msg-2", "file appeared"), - ev_completed("resp-4"), + ev_completed("resp-5"), ]), ) .await; - test.submit_turn("wait for resumed file").await?; + test.submit_turn_with_environments( + "wait for resumed file", + Some(vec![local(later_workspace.as_path().abs())]), + ) + .await?; let second_request = second_completion.single_request(); assert!( second_request - .function_call_output_text("call-2") + .function_call_output_text("call-3") .is_some_and(|output| output.ends_with("ready")) ); assert_eq!(fs::read_to_string(&resumed_file)?, "resumed"); + assert!(!later_resumed_file.exists()); Ok(()) } diff --git a/codex-rs/core/tests/suite/tools.rs b/codex-rs/core/tests/suite/tools.rs index 83584b8ec07..5cd926d7216 100644 --- a/codex-rs/core/tests/suite/tools.rs +++ b/codex-rs/core/tests/suite/tools.rs @@ -2,6 +2,7 @@ #![allow(clippy::unwrap_used)] use std::fs; +use std::sync::Arc; use std::time::Duration; use std::time::Instant; @@ -16,6 +17,9 @@ use codex_protocol::permissions::FileSystemSandboxEntry; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use core_test_support::PathExt; use core_test_support::assert_regex_match; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -28,11 +32,16 @@ use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::skip_if_sandbox; +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::local_selections; use core_test_support::test_codex::test_codex; use regex_lite::Regex; use serde_json::Value; use serde_json::json; +use tempfile::TempDir; +use tokio::sync::oneshot; fn tool_names(body: &Value) -> Vec { body.get("tools") @@ -130,6 +139,117 @@ async fn turn_environment_selection_keeps_environment_backed_tools() -> Result<( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sampling_request_keeps_one_environment_view_for_context_and_tool_execution() -> Result<()> +{ + skip_if_no_network!(Ok(())); + + let (release_tool_call_tx, release_tool_call_rx) = oneshot::channel(); + let (server, _completions) = start_streaming_sse_server(vec![ + vec![ + StreamingSseChunk { + gate: None, + body: sse(vec![ev_response_created("resp-1")]), + }, + StreamingSseChunk { + gate: Some(release_tool_call_rx), + body: sse(vec![ + ev_function_call( + "call-1", + "exec_command", + &serde_json::to_string(&json!({ + "cmd": "printf frozen > step-context-marker.txt", + }))?, + ), + ev_completed("resp-1"), + ]), + }, + ], + vec![StreamingSseChunk { + gate: None, + body: sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + }], + ]) + .await; + + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::UnifiedExec) + .expect("unified exec should enable for test"); + }); + let test = Arc::new(builder.build_with_streaming_server(&server).await?); + let initial_cwd = test.config.cwd.clone(); + let later_workspace = TempDir::new()?; + let later_cwd = later_workspace.path().abs(); + + let turn = { + let test = Arc::clone(&test); + let initial_cwd = initial_cwd.clone(); + tokio::spawn(async move { + test.submit_turn_with_environments( + "use the environment captured for this request", + Some(vec![local(initial_cwd)]), + ) + .await + }) + }; + server.wait_for_request_count(1).await; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "switch the live selection".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + environments: Some(local_selections(later_cwd.clone())), + ..Default::default() + }, + }) + .await?; + tokio::time::timeout(Duration::from_secs(2), async { + loop { + if test.codex.environment_selections().await == vec![local(later_cwd.clone())] { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .context("live environment selection did not update")?; + + release_tool_call_tx + .send(()) + .map_err(|_| anyhow::anyhow!("tool-call gate closed"))?; + turn.await??; + + let requests = server.requests().await; + assert_eq!(requests.len(), 2); + let first_request: Value = serde_json::from_slice(&requests[0])?; + assert!(tool_names(&first_request).contains(&"exec_command".to_string())); + assert!( + first_request["input"] + .to_string() + .contains(initial_cwd.as_path().to_string_lossy().as_ref()) + ); + assert_eq!( + fs::read_to_string(initial_cwd.join("step-context-marker.txt"))?, + "frozen" + ); + assert!(!later_cwd.join("step-context-marker.txt").exists()); + + server.shutdown().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn custom_tool_unknown_returns_custom_output_error() -> Result<()> { skip_if_no_network!(Ok(()));