fix: introduce EventMsg::TurnAborted

This commit is contained in:
Michael Bolin
2025-08-15 17:09:45 -07:00
parent 1ad8ae2579
commit 465175a163
8 changed files with 84 additions and 30 deletions

View File

@@ -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),
},

View File

@@ -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

View File

@@ -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<OutgoingMessageSender>,
codex_linux_sandbox_exe: Option<PathBuf>,
conversation_listeners: HashMap<Uuid, oneshot::Sender<()>>,
// Queue of pending interrupt requests per conversation. We reply when TurnAborted arrives.
pending_interrupts: Arc<Mutex<HashMap<Uuid, Vec<RequestId>>>>,
}
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<CodexConversation>,
outgoing: Arc<OutgoingMessageSender>,
pending_interrupts: Arc<Mutex<HashMap<Uuid, Vec<RequestId>>>>,
) {
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;
}
}
}
_ => {}
}
}

View File

@@ -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

View File

@@ -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(_))
}

View File

@@ -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")]

View File

@@ -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<String>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub enum TurnAbortReason {
Interrupted,
Replaced,
}
#[cfg(test)]
mod tests {
use super::*;

View File

@@ -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),