code-mode: use compact session-local cell ids

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();
 

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!(

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

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<Arc<dyn CodeModeSession>>,
     dispatch_broker: Arc<CodeModeDispatchBroker>,
+    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,

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),
     );