diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 30369f064c..139ae3985f 100644 --- a/codex-rs/Cargo.lock +++ b/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", diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index c669b01df3..b9ded54ae5 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/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 } diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index de367ec6a6..c05f39fbc6 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -220,6 +220,66 @@ impl HostState { } } +struct PendingDelegateCall { + cleanup: Option, +} + +struct PendingDelegateCallCleanup { + peer: Arc, + id: DelegateRequestId, + request_sent: bool, +} + +impl PendingDelegateCall { + fn new(peer: Arc, 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> { + 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, pending: Mutex>>>, @@ -240,7 +300,7 @@ impl HostPeer { } async fn call( - &self, + self: &Arc, 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) { if let Some(sender) = self.pending.lock().await.remove(&id) { let _ = sender.send(response); diff --git a/codex-rs/code-mode-host/tests/host.rs b/codex-rs/code-mode-host/tests/host.rs index ce31eb8a96..54be15088b 100644 --- a/codex-rs/code-mode-host/tests/host.rs +++ b/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}"); +}