code-mode: clean up canceled host delegates

2026-10-23 18:08:45 +00:00 · 2026-06-12 05:38:56 +00:00
parent df50f5cf11
commit 52402799db
4 changed files with 197 additions and 5 deletions
@@ -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}");
+}