code-mode: expose transport-neutral session runtime

codex-rs/code-mode-protocol/src/session.rs

@@ -63,9 +63,6 @@ pub trait CodeModeSessionDelegate: Send + Sync {
         text: String,
         cancellation_token: CancellationToken,
     ) -> NotificationFuture<'a>;
-
-    /// Releases delegate state associated with a cell after it reaches a terminal state.
-    fn cell_closed(&self, cell_id: &CellId);
 }
 
 /// A stateful code-mode session owned by one Codex thread.

codex-rs/code-mode/Cargo.toml

@@ -20,9 +20,11 @@ codex-code-mode-protocol = { workspace = true }
 codex-protocol = { workspace = true }
 deno_core_icudata = { workspace = true }
 serde_json = { workspace = true }
+sha2 = { workspace = true }
 tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }
 tokio-util = { workspace = true, features = ["rt"] }
 tracing = { workspace = true }
+uuid = { workspace = true, features = ["v4"] }
 v8 = { workspace = true }
 
 [dev-dependencies]

codex-rs/code-mode/src/cell_actor/callbacks.rs

@@ -1,4 +1,5 @@
 use std::sync::Arc;
+use std::time::Duration;
 
 use tokio::task::JoinSet;
 use tokio_util::sync::CancellationToken;
@@ -8,6 +9,8 @@ use super::CellHost;
 use super::CellToolCall;
 use crate::runtime::RuntimeCommand;
 
+const CALLBACK_CANCELLATION_GRACE: Duration = Duration::from_millis(100);
+
 #[derive(Clone, Copy)]
 pub(super) enum CallbackCompletion {
     DrainNotifications,
@@ -53,10 +56,13 @@ pub(super) async fn finish_callbacks(
 ) {
     if matches!(completion, CallbackCompletion::Cancel) {
         cancellation_token.cancel();
+        finish_cancelled_tasks(notification_tasks, "notification").await;
+        finish_cancelled_tasks(tool_tasks, "tool").await;
+        return;
     }
     drain_tasks(notification_tasks, "notification").await;
     cancellation_token.cancel();
-    drain_tasks(tool_tasks, "tool").await;
+    finish_cancelled_tasks(tool_tasks, "tool").await;
 }
 
 pub(super) fn log_task_result(
@@ -75,3 +81,17 @@ async fn drain_tasks(tasks: &mut JoinSet<()>, description: &str) {
         log_task_result(Some(result), description);
     }
 }
+
+async fn abort_tasks(tasks: &mut JoinSet<()>, description: &str) {
+    tasks.abort_all();
+    drain_tasks(tasks, description).await;
+}
+
+async fn finish_cancelled_tasks(tasks: &mut JoinSet<()>, description: &str) {
+    if tokio::time::timeout(CALLBACK_CANCELLATION_GRACE, drain_tasks(tasks, description))
+        .await
+        .is_err()
+    {
+        abort_tasks(tasks, description).await;
+    }
+}

codex-rs/code-mode/src/cell_actor/types.rs

@@ -20,7 +20,6 @@ use crate::session_runtime::ToolName;
 pub(crate) type CellEventFuture =
     Pin<Box<dyn Future<Output = Result<CellEvent, CellError>> + Send + 'static>>;
 
-#[allow(dead_code)]
 pub(crate) type ResumeFuture =
     Pin<Box<dyn Future<Output = Result<ResumeOutcome, CellError>> + Send + 'static>>;
 
@@ -122,7 +121,6 @@ impl CellHandle {
         response_event(response_rx)
     }
 
-    #[allow(dead_code)]
     pub(crate) fn resume(&self, generation: PendingGeneration) -> ResumeFuture {
         if !self.state.accepting_observations() {
             return closed_resume();
@@ -505,7 +503,6 @@ pub(super) enum CellCommand {
         mode: ObserveMode,
         response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
     },
-    #[allow(dead_code)]
     Resume {
         generation: PendingGeneration,
         response_tx: oneshot::Sender<Result<ResumeOutcome, CellError>>,
@@ -524,7 +521,6 @@ fn closed_event() -> CellEventFuture {
     Box::pin(async { Err(CellError::Closed) })
 }
 
-#[allow(dead_code)]
 fn closed_resume() -> ResumeFuture {
     Box::pin(async { Err(CellError::Closed) })
 }

codex-rs/code-mode/src/lib.rs

@@ -1,9 +1,10 @@
 mod cell_actor;
 mod runtime;
 mod service;
-mod session_runtime;
+pub mod session_runtime;
 
 pub use codex_code_mode_protocol::*;
 pub use service::CodeModeService;
 pub use service::InProcessCodeModeSessionProvider;
 pub use service::NoopCodeModeSessionDelegate;
+pub use session_runtime::SessionRuntime;

codex-rs/code-mode/src/service.rs

@@ -74,8 +74,6 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
     ) -> NotificationFuture<'a> {
         Box::pin(async { Ok(()) })
     }
-
-    fn cell_closed(&self, _cell_id: &CellId) {}
 }
 
 #[derive(Default)]
@@ -112,8 +110,14 @@ impl CodeModeService {
     }
 
     pub fn with_delegate(delegate: Arc<dyn CodeModeSessionDelegate>) -> Self {
+        #[cfg(not(test))]
+        let runtime = Arc::new(SessionRuntime::new(Arc::new(ProtocolDelegate { delegate })));
+        #[cfg(test)]
+        let runtime = Arc::new(SessionRuntime::new_for_test(Arc::new(ProtocolDelegate {
+            delegate,
+        })));
         Self {
-            runtime: Arc::new(SessionRuntime::new(Arc::new(ProtocolDelegate { delegate }))),
+            runtime,
             observations: Mutex::new(HashMap::new()),
             #[cfg(test)]
             pending_generations: Mutex::new(HashMap::new()),
@@ -179,7 +183,7 @@ impl CodeModeService {
             .ok_or_else(|| "observation ended before producing a result".to_string())?
     }
 
-    #[allow(dead_code)]
+    #[cfg(test)]
     async fn begin_observe(
         &self,
         request: ObserveRequest,
@@ -398,10 +402,6 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate {
             )
             .await
     }
-
-    fn cell_closed(&self, cell_id: &runtime::CellId) {
-        self.delegate.cell_closed(&protocol_cell_id(cell_id));
-    }
 }
 
 fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest {