code-mode: clean up canceled host delegates

codex-rs/Cargo.lock

@@ -2496,6 +2496,7 @@ version = "0.0.0"
 dependencies = [
  "codex-code-mode",
  "codex-code-mode-protocol",
+ "codex-protocol",
  "codex-utils-cargo-bin",
  "pretty_assertions",
  "tokio",

codex-rs/code-mode-host/Cargo.toml

@@ -14,13 +14,14 @@ workspace = true
 [dependencies]
 codex-code-mode = { workspace = true }
 codex-code-mode-protocol = { workspace = true }
-tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt-multi-thread", "sync"] }
+tokio = { workspace = true, features = ["io-std", "io-util", "macros", "process", "rt-multi-thread", "sync"] }
 tokio-util = { workspace = true, features = ["rt"] }
 # The host owns V8 process configuration, so enable the sandbox without an
 # internal workspace-crate feature gate.
 v8 = { workspace = true, features = ["v8_enable_sandbox"] }
 
 [dev-dependencies]
+codex-protocol = { workspace = true }
 codex-utils-cargo-bin = { workspace = true }
 pretty_assertions = { workspace = true }
 

codex-rs/code-mode-host/src/main.rs

@@ -220,6 +220,66 @@ impl HostState {
     }
 }
 
+struct PendingDelegateCall {
+    cleanup: Option<PendingDelegateCallCleanup>,
+}
+
+struct PendingDelegateCallCleanup {
+    peer: Arc<HostPeer>,
+    id: DelegateRequestId,
+    request_sent: bool,
+}
+
+impl PendingDelegateCall {
+    fn new(peer: Arc<HostPeer>, id: DelegateRequestId) -> Self {
+        Self {
+            cleanup: Some(PendingDelegateCallCleanup {
+                peer,
+                id,
+                request_sent: false,
+            }),
+        }
+    }
+
+    fn mark_request_sent(&mut self) {
+        if let Some(cleanup) = self.cleanup.as_mut() {
+            cleanup.request_sent = true;
+        }
+    }
+
+    fn disarm(&mut self) {
+        self.cleanup = None;
+    }
+
+    async fn cleanup(&mut self) {
+        if let Some(task) = self.spawn_cleanup() {
+            let _ = task.await;
+        }
+    }
+
+    fn spawn_cleanup(&mut self) -> Option<tokio::task::JoinHandle<()>> {
+        let handle = tokio::runtime::Handle::try_current().ok()?;
+        let cleanup = self.cleanup.take()?;
+        Some(handle.spawn(cleanup.run()))
+    }
+}
+
+impl PendingDelegateCallCleanup {
+    async fn run(self) {
+        if self.request_sent {
+            self.peer.cancel_pending(self.id).await;
+        } else {
+            self.peer.discard_pending(self.id).await;
+        }
+    }
+}
+
+impl Drop for PendingDelegateCall {
+    fn drop(&mut self) {
+        std::mem::drop(self.spawn_cleanup());
+    }
+}
+
 struct HostPeer {
     outgoing_tx: mpsc::Sender<HostMessage>,
     pending: Mutex<HashMap<DelegateRequestId, oneshot::Sender<Result<DelegateResponse, String>>>>,
@@ -240,7 +300,7 @@ impl HostPeer {
     }
 
     async fn call(
-        &self,
+        self: &Arc<Self>,
         session_id: SessionId,
         request: DelegateRequest,
         cancellation_token: CancellationToken,
@@ -248,6 +308,7 @@ impl HostPeer {
         let id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
         let (response_tx, response_rx) = oneshot::channel();
         self.pending.lock().await.insert(id, response_tx);
+        let mut pending_call = PendingDelegateCall::new(Arc::clone(self), id);
         if self
             .outgoing_tx
             .send(HostMessage::DelegateRequest {
@@ -258,21 +319,32 @@ impl HostPeer {
             .await
             .is_err()
         {
-            self.pending.lock().await.remove(&id);
+            pending_call.cleanup().await;
             return Err("code-mode client connection closed".to_string());
         }
+        pending_call.mark_request_sent();
         tokio::select! {
             response = response_rx => {
+                pending_call.disarm();
                 response.map_err(|_| "code-mode client closed before returning delegate output".to_string())?
             }
             _ = cancellation_token.cancelled() => {
-                self.pending.lock().await.remove(&id);
-                self.send(HostMessage::CancelDelegateRequest { id }).await;
+                pending_call.cleanup().await;
                 Err("code mode delegate request cancelled".to_string())
             }
         }
     }
 
+    async fn discard_pending(&self, id: DelegateRequestId) {
+        self.pending.lock().await.remove(&id);
+    }
+
+    async fn cancel_pending(&self, id: DelegateRequestId) {
+        if self.pending.lock().await.remove(&id).is_some() {
+            self.send(HostMessage::CancelDelegateRequest { id }).await;
+        }
+    }
+
     async fn complete(&self, id: DelegateRequestId, response: Result<DelegateResponse, String>) {
         if let Some(sender) = self.pending.lock().await.remove(&id) {
             let _ = sender.send(response);

codex-rs/code-mode-host/tests/host.rs

@@ -1,15 +1,18 @@
 use std::process::Stdio;
 use std::time::Duration;
 
+use codex_code_mode_protocol::CodeModeToolKind;
 use codex_code_mode_protocol::ExecuteRequest;
 use codex_code_mode_protocol::FunctionCallOutputContentItem;
 use codex_code_mode_protocol::RuntimeResponse;
+use codex_code_mode_protocol::ToolDefinition;
 use codex_code_mode_protocol::wire::ClientMessage;
 use codex_code_mode_protocol::wire::HostMessage;
 use codex_code_mode_protocol::wire::HostRequest;
 use codex_code_mode_protocol::wire::HostResponse;
 use codex_code_mode_protocol::wire::read_frame;
 use codex_code_mode_protocol::wire::write_frame;
+use codex_protocol::ToolName;
 use pretty_assertions::assert_eq;
 use tokio::process::Command;
 
@@ -123,3 +126,118 @@ async fn serves_code_mode_sessions_over_stdio() {
         .expect("wait for host");
     assert!(status.success(), "host exited with {status}");
 }
+
+#[tokio::test]
+async fn shutting_down_session_cancels_pending_delegate_request() {
+    let mut child = Command::new(
+        codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")
+            .expect("resolve codex-code-mode-host binary"),
+    )
+    .stdin(Stdio::piped())
+    .stdout(Stdio::piped())
+    .stderr(Stdio::inherit())
+    .kill_on_drop(true)
+    .spawn()
+    .expect("spawn codex-code-mode-host");
+    let mut stdin = child.stdin.take().expect("host stdin");
+    let mut stdout = child.stdout.take().expect("host stdout");
+
+    write_frame(
+        &mut stdin,
+        &ClientMessage::Request {
+            id: 1,
+            request: HostRequest::CreateSession,
+        },
+    )
+    .await
+    .expect("create session request");
+    let session_id = match read_frame(&mut stdout).await.expect("create response") {
+        Some(HostMessage::Response {
+            id: 1,
+            response: Ok(HostResponse::SessionCreated { session_id }),
+        }) => session_id,
+        message => panic!("unexpected create-session response: {message:?}"),
+    };
+
+    write_frame(
+        &mut stdin,
+        &ClientMessage::Request {
+            id: 2,
+            request: HostRequest::Execute {
+                session_id,
+                request: ExecuteRequest {
+                    tool_call_id: "call-1".to_string(),
+                    enabled_tools: vec![ToolDefinition {
+                        name: "echo".to_string(),
+                        tool_name: ToolName::plain("echo"),
+                        description: String::new(),
+                        kind: CodeModeToolKind::Function,
+                        input_schema: None,
+                        output_schema: None,
+                    }],
+                    source: "await tools.echo({});".to_string(),
+                    yield_time_ms: None,
+                    max_output_tokens: None,
+                },
+            },
+        },
+    )
+    .await
+    .expect("execute request");
+    let delegate_request_id = tokio::time::timeout(Duration::from_secs(5), async {
+        loop {
+            match read_frame(&mut stdout).await.expect("delegate request") {
+                Some(HostMessage::DelegateRequest {
+                    id,
+                    session_id: request_session_id,
+                    ..
+                }) if request_session_id == session_id => break id,
+                Some(HostMessage::Response {
+                    id: 2,
+                    response: Ok(HostResponse::ExecutionStarted { .. }),
+                }) => {}
+                message => panic!("unexpected execute message: {message:?}"),
+            }
+        }
+    })
+    .await
+    .expect("delegate request timeout");
+
+    write_frame(
+        &mut stdin,
+        &ClientMessage::Request {
+            id: 3,
+            request: HostRequest::ShutdownSession { session_id },
+        },
+    )
+    .await
+    .expect("shutdown request");
+    tokio::time::timeout(Duration::from_secs(5), async {
+        let mut cancellation_received = false;
+        let mut shutdown_received = false;
+        while !cancellation_received || !shutdown_received {
+            match read_frame(&mut stdout).await.expect("shutdown response") {
+                Some(HostMessage::CancelDelegateRequest { id }) => {
+                    assert_eq!(id, delegate_request_id);
+                    cancellation_received = true;
+                }
+                Some(HostMessage::Response {
+                    id: 3,
+                    response: Ok(HostResponse::SessionShutdown),
+                }) => shutdown_received = true,
+                Some(HostMessage::InitialResponse { .. })
+                | Some(HostMessage::CellClosed { .. }) => {}
+                message => panic!("unexpected shutdown message: {message:?}"),
+            }
+        }
+    })
+    .await
+    .expect("shutdown cancellation timeout");
+
+    drop(stdin);
+    let status = tokio::time::timeout(Duration::from_secs(5), child.wait())
+        .await
+        .expect("host exit timeout")
+        .expect("wait for host");
+    assert!(status.success(), "host exited with {status}");
+}