code-mode: preserve missing-cell wait results

codex-rs/core/src/tools/code_mode/mod.rs

@@ -14,6 +14,7 @@ use codex_code_mode_protocol::CodeModeSession;
 use codex_code_mode_protocol::CodeModeSessionProvider;
 use codex_code_mode_protocol::CodeModeToolKind;
 use codex_code_mode_protocol::RuntimeResponse;
+use codex_code_mode_protocol::WaitOutcome;
 use codex_protocol::models::FunctionCallOutputContentItem;
 use serde_json::Value as JsonValue;
 use std::sync::atomic::AtomicBool;
@@ -94,22 +95,20 @@ impl CodeModeService {
         &self,
         request: codex_code_mode_protocol::WaitRequest,
     ) -> Result<codex_code_mode_protocol::WaitOutcome, String> {
-        self.current_session()
-            .await
-            .ok_or_else(|| "code mode session is unavailable".to_string())?
-            .wait(request)
-            .await
+        match self.current_session().await {
+            Some(session) => session.wait(request).await,
+            None => Ok(missing_cell_outcome(request.cell_id)),
+        }
     }
 
     pub(crate) async fn terminate(
         &self,
         cell_id: CellId,
     ) -> Result<codex_code_mode_protocol::WaitOutcome, String> {
-        self.current_session()
-            .await
-            .ok_or_else(|| "code mode session is unavailable".to_string())?
-            .terminate(cell_id)
-            .await
+        match self.current_session().await {
+            Some(session) => session.terminate(cell_id).await,
+            None => Ok(missing_cell_outcome(cell_id)),
+        }
     }
 
     pub(crate) async fn shutdown(&self) -> Result<(), String> {
@@ -364,6 +363,14 @@ fn serialize_function_tool_arguments(
     }
 }
 
+fn missing_cell_outcome(cell_id: CellId) -> WaitOutcome {
+    WaitOutcome::MissingCell(RuntimeResponse::Result {
+        error_text: Some(format!("exec cell {cell_id} not found")),
+        cell_id,
+        content_items: Vec::new(),
+    })
+}
+
 fn build_freeform_tool_payload(
     tool_name: &ToolName,
     input: Option<JsonValue>,
@@ -536,20 +543,22 @@ mod tests {
                     cell_id: CellId::new("host1_1".to_string()),
                     yield_time_ms: 1,
                 })
-                .await
-                .err()
-                .as_deref(),
-            Some("code mode session is unavailable")
+                .await,
+            Ok(WaitOutcome::MissingCell(RuntimeResponse::Result {
+                cell_id: CellId::new("host1_1".to_string()),
+                content_items: Vec::new(),
+                error_text: Some("exec cell host1_1 not found".to_string()),
+            }))
         );
         assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1);
 
         assert_eq!(
-            service
-                .terminate(CellId::new("host1_1".to_string()))
-                .await
-                .err()
-                .as_deref(),
-            Some("code mode session is unavailable")
+            service.terminate(CellId::new("host1_1".to_string())).await,
+            Ok(WaitOutcome::MissingCell(RuntimeResponse::Result {
+                cell_id: CellId::new("host1_1".to_string()),
+                content_items: Vec::new(),
+                error_text: Some("exec cell host1_1 not found".to_string()),
+            }))
         );
         assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1);
 

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

@@ -1921,7 +1921,7 @@ text("after terminate");
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn code_mode_wait_returns_error_without_a_session() -> Result<()> {
+async fn code_mode_wait_returns_error_for_unknown_session() -> Result<()> {
     skip_if_no_network!(Ok(()));
 
     let server = responses::start_mock_server().await;
@@ -1964,10 +1964,17 @@ async fn code_mode_wait_returns_error_without_a_session() -> Result<()> {
     assert_ne!(success, Some(true));
 
     let items = function_tool_output_items(&request, "call-1");
-    assert_eq!(items.len(), 1);
-    assert_eq!(
+    assert_eq!(items.len(), 2);
+    assert_regex_match(
+        concat!(
+            r"(?s)\A",
+            r"Script failed\nWall time \d+\.\d seconds\nOutput:\n\z"
+        ),
         text_item(&items, /*index*/ 0),
-        "code mode session is unavailable"
+    );
+    assert_eq!(
+        text_item(&items, /*index*/ 1),
+        "Script error:\nexec cell 999999 not found"
     );
 
     Ok(())