From 465175a163f478b5682f560826d7dbcf29ef293a Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 15 Aug 2025 17:09:45 -0700 Subject: [PATCH] fix: introduce EventMsg::TurnAborted --- codex-rs/core/src/codex.rs | 24 ++++++------ .../src/event_processor_with_human_output.rs | 9 +++++ .../mcp-server/src/codex_message_processor.rs | 37 +++++++++++++++---- codex-rs/mcp-server/src/codex_tool_runner.rs | 1 + codex-rs/mcp-server/src/conversation_loop.rs | 25 +++++++++---- codex-rs/mcp-server/src/wire_format.rs | 7 +++- codex-rs/protocol/src/protocol.rs | 10 ++++- codex-rs/tui/src/chatwidget.rs | 1 + 8 files changed, 84 insertions(+), 30 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 020acd045e..e365afa1ee 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -14,6 +14,7 @@ use codex_apply_patch::ApplyPatchAction; use codex_apply_patch::MaybeApplyPatchVerified; use codex_apply_patch::maybe_parse_apply_patch_verified; use codex_login::CodexAuth; +use codex_protocol::protocol::TurnAbortReason; use futures::prelude::*; use mcp_types::CallToolResult; use serde::Serialize; @@ -535,7 +536,7 @@ impl Session { pub fn set_task(&self, task: AgentTask) { let mut state = self.state.lock_unchecked(); if let Some(current_task) = state.current_task.take() { - current_task.abort(); + current_task.abort(TurnAbortReason::Replaced); } state.current_task = Some(task); } @@ -852,13 +853,13 @@ impl Session { .await } - fn abort(&self) { - info!("Aborting existing session"); + fn interrupt_task(&self) { + info!("interrupt received: abort current task, if any"); let mut state = self.state.lock_unchecked(); state.pending_approvals.clear(); state.pending_input.clear(); if let Some(task) = state.current_task.take() { - task.abort(); + task.abort(TurnAbortReason::Interrupted); } } @@ -894,7 +895,7 @@ impl Session { impl Drop for Session { fn drop(&mut self) { - self.abort(); + self.interrupt_task(); } } @@ -964,14 +965,13 @@ impl AgentTask { } } - fn abort(self) { + fn abort(self, abort_reason: TurnAbortReason) { + // TOCTOU? if !self.handle.is_finished() { self.handle.abort(); let event = Event { id: self.sub_id, - msg: EventMsg::Error(ErrorEvent { - message: " Turn interrupted".to_string(), - }), + msg: EventMsg::TurnAborted(abort_reason), }; let tx_event = self.sess.tx_event.clone(); tokio::spawn(async move { @@ -994,7 +994,7 @@ async fn submission_loop( debug!(?sub, "Submission"); match sub.op { Op::Interrupt => { - sess.abort(); + sess.interrupt_task(); } Op::UserInput { items } => { // attempt to inject input into current task @@ -1065,13 +1065,13 @@ async fn submission_loop( } Op::ExecApproval { id, decision } => match decision { ReviewDecision::Abort => { - sess.abort(); + sess.interrupt_task(); } other => sess.notify_approval(&id, other), }, Op::PatchApproval { id, decision } => match decision { ReviewDecision::Abort => { - sess.abort(); + sess.interrupt_task(); } other => sess.notify_approval(&id, other), }, diff --git a/codex-rs/exec/src/event_processor_with_human_output.rs b/codex-rs/exec/src/event_processor_with_human_output.rs index 5d64fe620b..732345906c 100644 --- a/codex-rs/exec/src/event_processor_with_human_output.rs +++ b/codex-rs/exec/src/event_processor_with_human_output.rs @@ -21,6 +21,7 @@ use codex_core::protocol::PatchApplyBeginEvent; use codex_core::protocol::PatchApplyEndEvent; use codex_core::protocol::SessionConfiguredEvent; use codex_core::protocol::TaskCompleteEvent; +use codex_core::protocol::TurnAbortReason; use codex_core::protocol::TurnDiffEvent; use owo_colors::OwoColorize; use owo_colors::Style; @@ -522,6 +523,14 @@ impl EventProcessor for EventProcessorWithHumanOutput { EventMsg::GetHistoryEntryResponse(_) => { // Currently ignored in exec output. } + EventMsg::TurnAborted(abort_reason) => match abort_reason { + TurnAbortReason::Interrupted => { + ts_println!(self, "task interrupted"); + } + TurnAbortReason::Replaced => { + ts_println!(self, "task aborted: replaced by a new task"); + } + }, EventMsg::ShutdownComplete => return CodexStatus::Shutdown, } CodexStatus::Running diff --git a/codex-rs/mcp-server/src/codex_message_processor.rs b/codex-rs/mcp-server/src/codex_message_processor.rs index d930c03b71..fed77a408e 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -46,6 +46,7 @@ use crate::wire_format::SendUserTurnParams; use crate::wire_format::SendUserTurnResponse; use codex_core::protocol::InputItem as CoreInputItem; use codex_core::protocol::Op; +use tokio::sync::Mutex; /// Handles JSON-RPC messages for Codex conversations. pub(crate) struct CodexMessageProcessor { @@ -53,6 +54,8 @@ pub(crate) struct CodexMessageProcessor { outgoing: Arc, codex_linux_sandbox_exe: Option, conversation_listeners: HashMap>, + // Queue of pending interrupt requests per conversation. We reply when TurnAborted arrives. + pending_interrupts: Arc>>>, } impl CodexMessageProcessor { @@ -66,6 +69,7 @@ impl CodexMessageProcessor { outgoing, codex_linux_sandbox_exe, conversation_listeners: HashMap::new(), + pending_interrupts: Arc::new(Mutex::new(HashMap::new())), } } @@ -246,13 +250,14 @@ impl CodexMessageProcessor { return; }; - let _ = conversation.submit(Op::Interrupt).await; + // Record the pending interrupt so we can reply when TurnAborted arrives. + { + let mut map = self.pending_interrupts.lock().await; + map.entry(conversation_id.0).or_default().push(request_id); + } - // Apparently CodexConversation does not send an ack for Op::Interrupt, - // so we can reply to the request right away. - self.outgoing - .send_response(request_id, InterruptConversationResponse {}) - .await; + // Submit the interrupt; we'll respond upon TurnAborted. + let _ = conversation.submit(Op::Interrupt).await; } async fn add_conversation_listener( @@ -280,6 +285,7 @@ impl CodexMessageProcessor { self.conversation_listeners .insert(subscription_id, cancel_tx); let outgoing_for_task = self.outgoing.clone(); + let pending_interrupts = self.pending_interrupts.clone(); tokio::spawn(async move { loop { tokio::select! { @@ -320,7 +326,7 @@ impl CodexMessageProcessor { }) .await; - apply_bespoke_event_handling(event, conversation_id, conversation.clone(), outgoing_for_task.clone()).await; + apply_bespoke_event_handling(event.clone(), conversation_id, conversation.clone(), outgoing_for_task.clone(), pending_interrupts.clone()).await; } } } @@ -359,6 +365,7 @@ async fn apply_bespoke_event_handling( conversation_id: ConversationId, conversation: Arc, outgoing: Arc, + pending_interrupts: Arc>>>, ) { let Event { id: event_id, msg } = event; match msg { @@ -407,6 +414,22 @@ async fn apply_bespoke_event_handling( on_exec_approval_response(event_id, rx, conversation).await; }); } + // If this is a TurnAborted, reply to any pending interrupt requests. + EventMsg::TurnAborted(reason) => { + let pending = { + let mut map = pending_interrupts.lock().await; + map.remove(&conversation_id.0).unwrap_or_default() + }; + if !pending.is_empty() { + let response = InterruptConversationResponse { + abort_reason: reason, + }; + for rid in pending { + outgoing.send_response(rid, response.clone()).await; + } + } + } + _ => {} } } diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index ff660167d5..c0d14ecedc 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -272,6 +272,7 @@ async fn run_codex_tool_session_inner( | EventMsg::TurnDiff(_) | EventMsg::GetHistoryEntryResponse(_) | EventMsg::PlanUpdate(_) + | EventMsg::TurnAborted(_) | EventMsg::ShutdownComplete => { // For now, we do not do anything extra for these // events. Note that diff --git a/codex-rs/mcp-server/src/conversation_loop.rs b/codex-rs/mcp-server/src/conversation_loop.rs index 61f0c95ad4..278dd09cd0 100644 --- a/codex-rs/mcp-server/src/conversation_loop.rs +++ b/codex-rs/mcp-server/src/conversation_loop.rs @@ -7,6 +7,7 @@ use crate::patch_approval::handle_patch_approval_request; use codex_core::CodexConversation; use codex_core::protocol::AgentMessageEvent; use codex_core::protocol::ApplyPatchApprovalRequestEvent; +use codex_core::protocol::Event; use codex_core::protocol::EventMsg; use codex_core::protocol::ExecApprovalRequestEvent; use mcp_types::RequestId; @@ -27,12 +28,14 @@ pub async fn run_conversation_loop( loop { match codex.next_event().await { Ok(event) => { - outgoing - .send_event_as_notification( - &event, - Some(OutgoingNotificationMeta::new(Some(request_id.clone()))), - ) - .await; + if should_dispatch_notification_for_event(&event) { + outgoing + .send_event_as_notification( + &event, + Some(OutgoingNotificationMeta::new(Some(request_id.clone()))), + ) + .await; + } match event.msg { EventMsg::ExecApprovalRequest(ExecApprovalRequestEvent { @@ -52,7 +55,6 @@ pub async fn run_conversation_loop( call_id, ) .await; - continue; } EventMsg::Error(_) => { error!("Codex runtime error"); @@ -75,7 +77,6 @@ pub async fn run_conversation_loop( event.id.clone(), ) .await; - continue; } EventMsg::TaskComplete(_) => {} EventMsg::SessionConfigured(_) => { @@ -107,6 +108,7 @@ pub async fn run_conversation_loop( | EventMsg::PatchApplyEnd(_) | EventMsg::GetHistoryEntryResponse(_) | EventMsg::PlanUpdate(_) + | EventMsg::TurnAborted(_) | EventMsg::ShutdownComplete => { // For now, we do not do anything extra for these // events. Note that @@ -123,3 +125,10 @@ pub async fn run_conversation_loop( } } } + +fn should_dispatch_notification_for_event(event: &Event) -> bool { + // This should increase over time. Note we do not send a notification for + // TurnAborted because clients should look for the response to + // InterruptConversation instead. + !matches!(event.msg, EventMsg::TurnAborted(_)) +} diff --git a/codex-rs/mcp-server/src/wire_format.rs b/codex-rs/mcp-server/src/wire_format.rs index 2dca1b79b7..f7f6e30224 100644 --- a/codex-rs/mcp-server/src/wire_format.rs +++ b/codex-rs/mcp-server/src/wire_format.rs @@ -6,6 +6,7 @@ use codex_core::protocol::AskForApproval; use codex_core::protocol::FileChange; use codex_core::protocol::ReviewDecision; use codex_core::protocol::SandboxPolicy; +use codex_core::protocol::TurnAbortReason; use codex_core::protocol_config_types::ReasoningEffort; use codex_core::protocol_config_types::ReasoningSummary; use mcp_types::RequestId; @@ -152,9 +153,11 @@ pub struct InterruptConversationParams { pub conversation_id: ConversationId, } -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +#[derive(Serialize, Deserialize, Debug, Clone)] #[serde(rename_all = "camelCase")] -pub struct InterruptConversationResponse {} +pub struct InterruptConversationResponse { + pub abort_reason: TurnAbortReason, +} #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] #[serde(rename_all = "camelCase")] diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index c4f50b4ffa..bf2e4e3c58 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -39,7 +39,7 @@ pub struct Submission { #[non_exhaustive] pub enum Op { /// Abort current task. - /// This server sends no corresponding Event + /// This server sends [`EventMsg::TurnAborted`] in response. Interrupt, /// Input from the user @@ -422,6 +422,8 @@ pub enum EventMsg { PlanUpdate(UpdatePlanArgs), + TurnAborted(TurnAbortReason), + /// Notification that the agent is shutting down. ShutdownComplete, } @@ -745,6 +747,12 @@ pub struct Chunk { pub inserted_lines: Vec, } +#[derive(Debug, Clone, Deserialize, Serialize)] +pub enum TurnAbortReason { + Interrupted, + Replaced, +} + #[cfg(test)] mod tests { use super::*; diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 637d50bd33..6b8207fbeb 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -631,6 +631,7 @@ impl ChatWidget<'_> { EventMsg::TaskComplete(TaskCompleteEvent { .. }) => self.on_task_complete(), EventMsg::TokenCount(token_usage) => self.on_token_count(token_usage), EventMsg::Error(ErrorEvent { message }) => self.on_error(message), + EventMsg::TurnAborted(_) => self.on_error("Turn interrupted".to_owned()), EventMsg::PlanUpdate(update) => self.on_plan_update(update), EventMsg::ExecApprovalRequest(ev) => self.on_exec_approval_request(id, ev), EventMsg::ApplyPatchApprovalRequest(ev) => self.on_apply_patch_approval_request(id, ev),