diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs index 6d6c0004d3..baa0704b01 100644 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -483,11 +483,13 @@ async fn natural_completion_waits_for_notifications_before_responding() { next_event(&mut events_rx).await, DelegateEvent::NotificationStarted ); - assert!( - tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut initial_response) - .await - .is_err() - ); + std::future::poll_fn(|context| match initial_response.as_mut().poll(context) { + std::task::Poll::Pending => std::task::Poll::Ready(()), + std::task::Poll::Ready(result) => { + panic!("execute returned while the notification was blocked: {result:?}") + } + }) + .await; delegate.release_notification(); diff --git a/codex-rs/code-mode/src/service_tests.rs b/codex-rs/code-mode/src/service_tests.rs index 4cc0066e3a..14604b38aa 100644 --- a/codex-rs/code-mode/src/service_tests.rs +++ b/codex-rs/code-mode/src/service_tests.rs @@ -354,16 +354,27 @@ await Promise.all([ } ); - tokio::time::sleep(Duration::from_millis(1100)).await; - - let resumed_response = tokio::time::timeout( - Duration::from_secs(1), - service.wait_to_pending(WaitToPendingRequest { - cell_id: cell_id("1"), - }), - ) + let resumed_response = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let response = service + .wait_to_pending(WaitToPendingRequest { + cell_id: cell_id("1"), + }) + .await + .unwrap(); + if matches!( + &response, + WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { + pending_tool_call_ids, + .. + }) if !pending_tool_call_ids.is_empty() + ) { + break response; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) .await - .unwrap() .unwrap(); assert_eq!( 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 904a29214b..8bb7b00ab8 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -39,7 +39,7 @@ impl CodeModeExecuteHandler { let enabled_tools = codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); let started_at = std::time::Instant::now(); - let cell_id = codex_code_mode::CellId::new(uuid::Uuid::new_v4().to_string()); + let cell_id = exec.session.services.code_mode_service.allocate_cell_id(); let runtime_cell_id = cell_id.to_string(); let code_cell_trace = exec .session diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 43258b0244..c18459746b 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -6,6 +6,8 @@ mod wait_handler; pub(crate) mod wait_spec; use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use std::time::Duration; use codex_code_mode::CellId; @@ -61,6 +63,7 @@ pub(crate) struct ExecContext { pub(crate) struct CodeModeService { session: Option>, dispatch_broker: Arc, + next_cell_id: AtomicU64, } impl CodeModeService { @@ -71,9 +74,18 @@ impl CodeModeService { dispatch_broker.clone(), ))), dispatch_broker, + next_cell_id: AtomicU64::new(1), } } + pub(crate) fn allocate_cell_id(&self) -> CellId { + CellId::new( + self.next_cell_id + .fetch_add(1, Ordering::Relaxed) + .to_string(), + ) + } + pub(crate) async fn execute( &self, request: codex_code_mode::ExecuteRequest, diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 10a8c4eb9c..60b476d614 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1354,7 +1354,7 @@ text("phase 3"); assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID [^\n]+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&first_items, /*index*/ 0), ); @@ -1395,7 +1395,7 @@ text("phase 3"); assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID [^\n]+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&second_items, /*index*/ 0), ); @@ -1498,7 +1498,7 @@ while (true) {} assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID [^\n]+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&first_items, /*index*/ 0), ); @@ -2177,7 +2177,7 @@ text("session b done"); assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID [^\n]+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&third_items, /*index*/ 0), ); @@ -2299,7 +2299,7 @@ text("after yield"); assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID [^\n]+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&first_items, /*index*/ 0), );