[codex] add code-mode host failure supervision hooks

codex-rs/code-mode/Cargo.toml

@@ -19,6 +19,7 @@ workspace = true
 codex-code-mode-protocol = { workspace = true }
 codex-protocol = { workspace = true }
 deno_core_icudata = { workspace = true }
+futures = { workspace = true }
 serde_json = { workspace = true }
 tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }
 tokio-util = { workspace = true, features = ["rt"] }

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

@@ -1,11 +1,14 @@
+use std::panic::AssertUnwindSafe;
 use std::sync::Arc;
 
+use futures::FutureExt;
 use tokio::task::JoinSet;
 use tokio_util::sync::CancellationToken;
 use tracing::warn;
 
 use super::CellHost;
 use super::CellToolCall;
+use crate::TaskFailureHandler;
 use crate::runtime::RuntimeCommand;
 
 #[derive(Clone, Copy)]
@@ -20,10 +23,20 @@ pub(super) fn spawn_notification<H: CellHost>(
     call_id: String,
     text: String,
     cancellation_token: CancellationToken,
+    task_failure_handler: Option<TaskFailureHandler>,
 ) {
     tasks.spawn(async move {
-        if let Err(err) = host.notify(call_id, text, cancellation_token).await {
-            warn!("failed to deliver code mode notification: {err}");
+        let callback =
+            AssertUnwindSafe(async move { host.notify(call_id, text, cancellation_token).await })
+                .catch_unwind()
+                .await;
+        match callback {
+            Ok(Ok(())) => {}
+            Ok(Err(err)) => warn!("failed to deliver code mode notification: {err}"),
+            Err(_) => report_task_failure(
+                task_failure_handler.as_ref(),
+                "code mode notification task panicked".to_string(),
+            ),
         }
     });
 }
@@ -34,14 +47,32 @@ pub(super) fn spawn_tool<H: CellHost>(
     invocation: CellToolCall,
     runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
     cancellation_token: CancellationToken,
+    task_failure_handler: Option<TaskFailureHandler>,
 ) {
     tasks.spawn(async move {
         let id = invocation.id.clone();
-        let command = match host.invoke_tool(invocation, cancellation_token).await {
-            Ok(result) => RuntimeCommand::ToolResponse { id, result },
-            Err(error_text) => RuntimeCommand::ToolError { id, error_text },
+        let callback =
+            AssertUnwindSafe(async move { host.invoke_tool(invocation, cancellation_token).await })
+                .catch_unwind()
+                .await;
+        let (command, failure_reason) = match callback {
+            Ok(Ok(result)) => (RuntimeCommand::ToolResponse { id, result }, None),
+            Ok(Err(error_text)) => (RuntimeCommand::ToolError { id, error_text }, None),
+            Err(_) => {
+                let failure_reason = "code mode tool task panicked".to_string();
+                (
+                    RuntimeCommand::ToolError {
+                        id,
+                        error_text: failure_reason.clone(),
+                    },
+                    Some(failure_reason),
+                )
+            }
         };
         let _ = runtime_tx.send(command);
+        if let Some(failure_reason) = failure_reason {
+            report_task_failure(task_failure_handler.as_ref(), failure_reason);
+        }
     });
 }
 
@@ -50,28 +81,48 @@ pub(super) async fn finish_callbacks(
     notification_tasks: &mut JoinSet<()>,
     tool_tasks: &mut JoinSet<()>,
     completion: CallbackCompletion,
+    task_failure_handler: Option<&TaskFailureHandler>,
 ) {
     if matches!(completion, CallbackCompletion::Cancel) {
         cancellation_token.cancel();
     }
-    drain_tasks(notification_tasks, "notification").await;
+    drain_tasks(notification_tasks, "notification", task_failure_handler).await;
     cancellation_token.cancel();
-    drain_tasks(tool_tasks, "tool").await;
+    drain_tasks(tool_tasks, "tool", task_failure_handler).await;
 }
 
-pub(super) fn log_task_result(
+pub(super) fn report_task_result(
     task_result: Option<Result<(), tokio::task::JoinError>>,
     description: &str,
+    task_failure_handler: Option<&TaskFailureHandler>,
 ) {
     if let Some(Err(err)) = task_result
         && !err.is_cancelled()
     {
-        warn!("code mode {description} task failed: {err}");
+        report_task_failure(
+            task_failure_handler,
+            format!("code mode {description} task failed: {err}"),
+        );
+    }
+}
+
+fn report_task_failure(task_failure_handler: Option<&TaskFailureHandler>, failure_reason: String) {
+    warn!("{failure_reason}");
+    if let Some(task_failure_handler) = task_failure_handler {
+        task_failure_handler(failure_reason);
     }
 }
 
-async fn drain_tasks(tasks: &mut JoinSet<()>, description: &str) {
+async fn drain_tasks(
+    tasks: &mut JoinSet<()>,
+    description: &str,
+    task_failure_handler: Option<&TaskFailureHandler>,
+) {
     while let Some(result) = tasks.join_next().await {
-        log_task_result(Some(result), description);
+        report_task_result(Some(result), description, task_failure_handler);
     }
 }
+
+#[cfg(test)]
+#[path = "callbacks_tests.rs"]
+mod tests;

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

@@ -0,0 +1,132 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::mpsc as std_mpsc;
+use std::time::Duration;
+
+use pretty_assertions::assert_eq;
+use serde_json::Value as JsonValue;
+use tokio::sync::mpsc;
+use tokio::task::JoinSet;
+use tokio_util::sync::CancellationToken;
+
+use super::*;
+use crate::cell_actor::CellState;
+use crate::cell_actor::CompletionCommit;
+use crate::runtime::RuntimeCommand;
+use crate::session_runtime::CellEvent;
+use crate::session_runtime::ToolKind;
+use crate::session_runtime::ToolName;
+
+struct PanickingCallbackHost;
+
+impl CellHost for PanickingCallbackHost {
+    async fn invoke_tool(
+        &self,
+        _invocation: CellToolCall,
+        _cancellation_token: CancellationToken,
+    ) -> Result<JsonValue, String> {
+        panic!("tool callback panic probe");
+    }
+
+    async fn notify(
+        &self,
+        _call_id: String,
+        _text: String,
+        _cancellation_token: CancellationToken,
+    ) -> Result<(), String> {
+        panic!("notification callback panic probe");
+    }
+
+    async fn commit_completion(
+        &self,
+        _stored_value_writes: HashMap<String, JsonValue>,
+        _event: CellEvent,
+        _pending_initial_yield_items: Option<Vec<crate::session_runtime::OutputItem>>,
+        _cell_state: Arc<CellState>,
+    ) -> CompletionCommit {
+        panic!("unexpected completion commit");
+    }
+
+    async fn closed(&self) {}
+}
+
+#[tokio::test]
+async fn tool_callback_panic_rejects_the_js_promise_and_reports_failure() {
+    let mut tasks = JoinSet::new();
+    let (runtime_tx, runtime_rx) = std_mpsc::channel();
+    let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
+    spawn_tool(
+        &mut tasks,
+        Arc::new(PanickingCallbackHost),
+        CellToolCall {
+            id: "tool-1".to_string(),
+            name: ToolName {
+                name: "panic".to_string(),
+                namespace: None,
+            },
+            kind: ToolKind::Function,
+            input: None,
+        },
+        runtime_tx,
+        CancellationToken::new(),
+        Some(Arc::new(move |reason| {
+            let _ = failure_tx.send(reason);
+        })),
+    );
+
+    tasks
+        .join_next()
+        .await
+        .expect("tool callback task")
+        .expect("tool callback wrapper");
+    let command = runtime_rx
+        .recv_timeout(Duration::from_secs(1))
+        .expect("tool error command");
+    let RuntimeCommand::ToolError { id, error_text } = command else {
+        panic!("expected a tool error command");
+    };
+    assert_eq!(id, "tool-1");
+    assert_eq!(error_text, "code mode tool task panicked");
+    assert_eq!(failure_rx.recv().await, Some(error_text));
+}
+
+#[tokio::test]
+async fn notification_callback_panic_reports_failure() {
+    let mut tasks = JoinSet::new();
+    let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
+    spawn_notification(
+        &mut tasks,
+        Arc::new(PanickingCallbackHost),
+        "notify-1".to_string(),
+        "hello".to_string(),
+        CancellationToken::new(),
+        Some(Arc::new(move |reason| {
+            let _ = failure_tx.send(reason);
+        })),
+    );
+
+    tasks
+        .join_next()
+        .await
+        .expect("notification callback task")
+        .expect("notification callback wrapper");
+    let failure_reason = failure_rx.recv().await.expect("notification failure");
+    assert_eq!(failure_reason, "code mode notification task panicked");
+}
+
+#[tokio::test]
+async fn callback_wrapper_join_error_reports_failure() {
+    let task_result = tokio::spawn(async {
+        panic!("callback wrapper panic probe");
+    })
+    .await;
+    let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
+    let task_failure_handler: TaskFailureHandler = Arc::new(move |reason| {
+        let _ = failure_tx.send(reason);
+    });
+
+    report_task_result(Some(task_result), "tool", Some(&task_failure_handler));
+
+    let failure_reason = failure_rx.recv().await.expect("wrapper failure");
+    assert!(failure_reason.contains("code mode tool task failed"));
+}