diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 9224ba45a4..aa0e064306 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/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 { - 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 { - 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, @@ -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); diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 11d78986b9..0f954bac7b 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/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(())