Stabilize remote routing e2e tests

codex-rs/core/tests/common/test_codex.rs

@@ -296,8 +296,11 @@ impl TestCodexBuilder {
         };
         let base_url = format!("{}/v1", server.uri());
         let test_env = TestEnv::local().await?;
-        Box::pin(self.build_with_home_and_base_url(base_url, home, /*resume_from*/ None, test_env))
-            .await
+        Box::pin(self.build_with_home_and_base_url(
+            base_url, home, /*resume_from*/ None, test_env,
+            /*include_local_for_remote*/ false,
+        ))
+        .await
     }
 
     pub async fn build_remote_aware(
@@ -310,8 +313,28 @@ impl TestCodexBuilder {
         };
         let base_url = format!("{}/v1", server.uri());
         let test_env = test_env().await?;
-        Box::pin(self.build_with_home_and_base_url(base_url, home, /*resume_from*/ None, test_env))
-            .await
+        Box::pin(self.build_with_home_and_base_url(
+            base_url, home, /*resume_from*/ None, test_env,
+            /*include_local_for_remote*/ false,
+        ))
+        .await
+    }
+
+    pub async fn build_remote_and_local_aware(
+        &mut self,
+        server: &wiremock::MockServer,
+    ) -> anyhow::Result<TestCodex> {
+        let home = match self.home.clone() {
+            Some(home) => home,
+            None => Arc::new(TempDir::new()?),
+        };
+        let base_url = format!("{}/v1", server.uri());
+        let test_env = test_env().await?;
+        Box::pin(self.build_with_home_and_base_url(
+            base_url, home, /*resume_from*/ None, test_env,
+            /*include_local_for_remote*/ true,
+        ))
+        .await
     }
 
     pub async fn build_with_streaming_server(
@@ -329,6 +352,7 @@ impl TestCodexBuilder {
             home,
             /*resume_from*/ None,
             test_env,
+            /*include_local_for_remote*/ false,
         ))
         .await
     }
@@ -350,8 +374,11 @@ impl TestCodexBuilder {
             config.realtime.version = RealtimeWsVersion::V1;
         }));
         let test_env = TestEnv::local().await?;
-        Box::pin(self.build_with_home_and_base_url(base_url, home, /*resume_from*/ None, test_env))
-            .await
+        Box::pin(self.build_with_home_and_base_url(
+            base_url, home, /*resume_from*/ None, test_env,
+            /*include_local_for_remote*/ false,
+        ))
+        .await
     }
 
     pub async fn resume(
@@ -362,8 +389,14 @@ impl TestCodexBuilder {
     ) -> anyhow::Result<TestCodex> {
         let base_url = format!("{}/v1", server.uri());
         let test_env = TestEnv::local().await?;
-        Box::pin(self.build_with_home_and_base_url(base_url, home, Some(rollout_path), test_env))
-            .await
+        Box::pin(self.build_with_home_and_base_url(
+            base_url,
+            home,
+            Some(rollout_path),
+            test_env,
+            /*include_local_for_remote*/ false,
+        ))
+        .await
     }
 
     async fn build_with_home_and_base_url(
@@ -372,6 +405,7 @@ impl TestCodexBuilder {
         home: Arc<TempDir>,
         resume_from: Option<PathBuf>,
         test_env: TestEnv,
+        include_local_for_remote: bool,
     ) -> anyhow::Result<TestCodex> {
         let (config, fallback_cwd) = self
             .prepare_config(base_url, &home, test_env.cwd().clone())
@@ -391,13 +425,30 @@ impl TestCodexBuilder {
             std::env::current_exe()?,
             codex_linux_sandbox_exe,
         )?;
-        let environment_manager = Arc::new(
-            codex_exec_server::EnvironmentManager::create_for_tests(
-                exec_server_url,
-                local_runtime_paths,
-            )
-            .await,
-        );
+        let environment_manager = match (exec_server_url, include_local_for_remote) {
+            (Some(exec_server_url), true) => {
+                codex_exec_server::EnvironmentManager::create_remote_aware_for_tests(
+                    exec_server_url,
+                    local_runtime_paths,
+                )
+                .await?
+            }
+            (Some(exec_server_url), false) => {
+                codex_exec_server::EnvironmentManager::create_for_tests(
+                    Some(exec_server_url),
+                    local_runtime_paths,
+                )
+                .await
+            }
+            (None, _) => {
+                codex_exec_server::EnvironmentManager::create_for_tests(
+                    /*exec_server_url*/ None,
+                    local_runtime_paths,
+                )
+                .await
+            }
+        };
+        let environment_manager = Arc::new(environment_manager);
         let file_system = test_env.environment().get_filesystem();
         let mut workspace_setups = vec![];
         swap(&mut self.workspace_setups, &mut workspace_setups);

codex-rs/core/tests/suite/remote_env.rs

@@ -58,7 +58,7 @@ async fn unified_exec_test(server: &wiremock::MockServer) -> Result<TestCodex> {
             "unified exec should enable for test: {result:?}",
         );
     });
-    builder.build_remote_aware(server).await
+    builder.build_remote_and_local_aware(server).await
 }
 
 async fn submit_turn_with_approval_and_environments(

codex-rs/core/tests/suite/view_image.rs

@@ -37,7 +37,6 @@ use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_function_call;
 use core_test_support::responses::ev_response_created;
 use core_test_support::responses::mount_models_once;
-use core_test_support::responses::mount_sse_sequence;
 use core_test_support::responses::sse;
 use core_test_support::responses::start_mock_server;
 use core_test_support::skip_if_no_network;
@@ -422,7 +421,7 @@ async fn view_image_routes_to_selected_local_environment() -> anyhow::Result<()>
     )
     .await?;
     let call_id = "call-view-image-local-env";
-    let response_mock = mount_sse_sequence(
+    let response_mock = responses::mount_sse_sequence(
         &server,
         vec![
             sse(vec![
@@ -493,7 +492,7 @@ async fn view_image_tool_applies_local_sandbox_read_denies() -> anyhow::Result<(
     )
     .await?;
     let call_id = "call-view-image-outside-cwd";
-    let response_mock = mount_sse_sequence(
+    let response_mock = responses::mount_sse_sequence(
         &server,
         vec![
             sse(vec![
@@ -562,7 +561,7 @@ async fn view_image_routes_to_selected_remote_environment() -> anyhow::Result<()
 
     let server = start_mock_server().await;
     let mut builder = test_codex();
-    let test = builder.build_remote_aware(&server).await?;
+    let test = builder.build_remote_and_local_aware(&server).await?;
     let local_cwd = TempDir::new()?;
     fs::write(local_cwd.path().join("remote.png"), b"not a remote image")?;
     let local_selection = TurnEnvironmentSelection {
@@ -582,18 +581,19 @@ async fn view_image_routes_to_selected_remote_environment() -> anyhow::Result<()
             /*sandbox*/ None,
         )
         .await?;
-    let png = BASE64_STANDARD.decode(
-        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII=",
-    )?;
     test.fs()
-        .write_file(&image_path, png, /*sandbox*/ None)
+        .write_file(
+            &image_path,
+            png_bytes(/*width*/ 1, /*height*/ 1, [255, 0, 0, 255])?,
+            /*sandbox*/ None,
+        )
         .await?;
     let remote_selection = TurnEnvironmentSelection {
         environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
         cwd: remote_cwd.clone(),
     };
     let call_id = "call-view-image-multi-env";
-    let response_mock = mount_sse_sequence(
+    let response_mock = responses::mount_sse_sequence(
         &server,
         vec![
             sse(vec![
@@ -626,9 +626,8 @@ async fn view_image_routes_to_selected_remote_environment() -> anyhow::Result<()
 
     let output = response_mock
         .last_request()
-        .context("missing request containing view_image output")?
-        .function_call_output(call_id)
-        .clone();
+        .context("missing request containing remote view_image output")?
+        .function_call_output(call_id);
     let output_items = output
         .get("output")
         .and_then(Value::as_array)

codex-rs/exec-server/src/environment.rs

@@ -91,6 +91,23 @@ impl EnvironmentManager {
         Self::from_default_provider_url(exec_server_url, local_runtime_paths).await
     }
 
+    /// Builds a test-only manager that keeps the local environment addressable
+    /// alongside a remote default environment.
+    pub async fn create_remote_aware_for_tests(
+        exec_server_url: String,
+        local_runtime_paths: ExecServerRuntimePaths,
+    ) -> Result<Self, ExecServerError> {
+        let snapshot = EnvironmentProviderSnapshot {
+            environments: vec![(
+                REMOTE_ENVIRONMENT_ID.to_string(),
+                Environment::remote_inner(exec_server_url, /*local_runtime_paths*/ None),
+            )],
+            default: EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()),
+            include_local: true,
+        };
+        Self::from_provider_snapshot(snapshot, local_runtime_paths)
+    }
+
     /// Builds a manager from `CODEX_EXEC_SERVER_URL` and local runtime paths
     /// used when creating local filesystem helpers.
     pub async fn new(args: EnvironmentManagerArgs) -> Self {
@@ -528,6 +545,34 @@ mod tests {
         assert!(!manager.local_environment().is_remote());
     }
 
+    #[tokio::test]
+    async fn remote_aware_test_manager_keeps_local_environment_addressable()
+    -> Result<(), ExecServerError> {
+        let manager = EnvironmentManager::create_remote_aware_for_tests(
+            "ws://127.0.0.1:8765".to_string(),
+            test_runtime_paths(),
+        )
+        .await?;
+
+        assert_eq!(
+            manager.default_environment_id(),
+            Some(REMOTE_ENVIRONMENT_ID)
+        );
+        assert!(
+            manager
+                .get_environment(REMOTE_ENVIRONMENT_ID)
+                .expect("remote environment")
+                .is_remote()
+        );
+        assert!(
+            !manager
+                .get_environment(LOCAL_ENVIRONMENT_ID)
+                .expect("local environment")
+                .is_remote()
+        );
+        Ok(())
+    }
+
     #[tokio::test]
     async fn environment_manager_default_environment_caches_environment() {
         let manager = EnvironmentManager::default_for_tests();