From 486264137c278504f4c5b3afd23d9f7f331bbb87 Mon Sep 17 00:00:00 2001
From: starr-openai <starr@openai.com>
Date: Sat, 9 May 2026 14:00:07 -0700
Subject: [PATCH] Fix remote-aware harness followups

Return remote-aware environment construction errors instead of using expect in library code, and accept the structured view_image function-call output shape captured by the response mock.

Co-authored-by: Codex <noreply@openai.com>
---
 codex-rs/core/tests/common/test_codex.rs | 7 ++++---
 codex-rs/core/tests/suite/view_image.rs  | 6 +++---
 codex-rs/exec-server/src/environment.rs  | 9 +++++----
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs
index 1e771db967..a374c55c06 100644
--- a/codex-rs/core/tests/common/test_codex.rs
+++ b/codex-rs/core/tests/common/test_codex.rs
@@ -386,13 +386,13 @@ impl TestCodexBuilder {
             std::env::current_exe()?,
             /*codex_linux_sandbox_exe*/ None,
         )?;
-        let environment_manager = Arc::new(match exec_server_url {
+        let environment_manager = match exec_server_url {
             Some(exec_server_url) => {
                 codex_exec_server::EnvironmentManager::create_remote_aware_for_tests(
                     exec_server_url,
                     local_runtime_paths,
                 )
-                .await
+                .await?
             }
             None => {
                 codex_exec_server::EnvironmentManager::create_for_tests(
@@ -401,7 +401,8 @@ impl TestCodexBuilder {
                 )
                 .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);
diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs
index badbe2ba4f..29b754d973 100644
--- a/codex-rs/core/tests/suite/view_image.rs
+++ b/codex-rs/core/tests/suite/view_image.rs
@@ -579,9 +579,9 @@ async fn view_image_routes_to_selected_remote_environment() -> anyhow::Result<()
 
     let output = wait_for_function_call_output(&test, &response_mock, call_id).await?;
     assert_eq!(response_mock.requests().len(), 2);
-    let output_items = output
-        .get("output")
-        .and_then(Value::as_array)
+    let output_value = output.get("output").unwrap_or(&output);
+    let output_items = output_value
+        .as_array()
         .context("view_image output should be content items")?;
     assert_eq!(output_items.len(), 1);
     let image_url = output_items[0]
diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs
index 703c8266e2..bcdd62c62e 100644
--- a/codex-rs/exec-server/src/environment.rs
+++ b/codex-rs/exec-server/src/environment.rs
@@ -95,7 +95,7 @@ impl EnvironmentManager {
     pub async fn create_remote_aware_for_tests(
         exec_server_url: String,
         local_runtime_paths: ExecServerRuntimePaths,
-    ) -> Self {
+    ) -> Result<Self, ExecServerError> {
         let snapshot = EnvironmentProviderSnapshot {
             environments: vec![(
                 REMOTE_ENVIRONMENT_ID.to_string(),
@@ -105,7 +105,6 @@ impl EnvironmentManager {
             include_local: true,
         };
         Self::from_provider_snapshot(snapshot, local_runtime_paths)
-            .expect("remote-aware test provider should create valid environments")
     }
 
     /// Builds a manager from `CODEX_EXEC_SERVER_URL` and local runtime paths
@@ -503,12 +502,13 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn remote_aware_test_manager_keeps_local_environment_addressable() {
+    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;
+        .await?;
 
         assert_eq!(
             manager.default_environment_id(),
@@ -526,6 +526,7 @@ mod tests {
                 .expect("local environment")
                 .is_remote()
         );
+        Ok(())
     }
 
     #[tokio::test]