openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-13 21:30:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(core): preserve early interrupted turns

Browse Source
This commit is contained in:
Felipe Coury
2026-07-09 21:31:54 -03:00
parent 4795329ff3
commit ba409a3fae
9 changed files with 593 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/hook_runtime.rs
+21 -8
View File
@@ -53,6 +53,12 @@ pub(crate) struct HookRuntimeOutcome {
pub additional_contexts: Vec<String>,
}
#[derive(Clone, Copy)]
pub(crate) enum UserInputEventDisposition {
Emit,
AlreadyEmitted,
}
pub(crate) enum PreToolUseHookResult {
Continue { updated_input: Option<Value> },
Blocked(String),
@@ -541,16 +547,23 @@ pub(crate) async fn record_pending_input(
turn_context: &Arc<TurnContext>,
pending_input: TurnInput,
additional_contexts: Vec<String>,
user_input_event_disposition: UserInputEventDisposition,
) {
match pending_input {
TurnInput::UserInput { content, client_id } => {
sess.record_user_prompt_and_emit_turn_item(
turn_context.as_ref(),
content.as_slice(),
client_id,
)
.await;
}
TurnInput::UserInput { content, client_id } => match user_input_event_disposition {
UserInputEventDisposition::Emit => {
sess.record_user_prompt_and_emit_turn_item(
turn_context.as_ref(),
content.as_slice(),
client_id,
)
.await;
}
UserInputEventDisposition::AlreadyEmitted => {
sess.record_user_prompt(turn_context.as_ref(), content.as_slice())
.await;
}
},
TurnInput::ResponseItem(item) => {
sess.record_conversation_items(turn_context, std::slice::from_ref(&item))
.await;
codex-rs/core/src/session/mod.rs
+42 -1
View File
@@ -3838,12 +3838,53 @@ impl Session {
let response_item = self.response_item_from_user_input(input.to_vec());
self.record_conversation_items(turn_context, std::slice::from_ref(&response_item))
.await;
self.emit_user_input_turn_item(turn_context, input, client_id)
.await;
self.ensure_rollout_materialized().await;
}
pub(crate) async fn record_user_prompt(&self, turn_context: &TurnContext, input: &[UserInput]) {
let response_item = self.response_item_from_user_input(input.to_vec());
self.record_conversation_items(turn_context, std::slice::from_ref(&response_item))
.await;
self.ensure_rollout_materialized().await;
}
/// Persist transcript lifecycle events for a regular turn's initial user input.
///
/// Accepted prompts use this after submission hooks. The interrupt fallback uses it with the
/// original input when cancellation wins before those hooks reach the persistence boundary.
/// This method intentionally does not add the prompt to model-visible history.
pub(crate) async fn record_initial_user_input_events(
&self,
turn_context: &TurnContext,
input: &[TurnInput],
) {
let mut emitted = false;
for input in input {
let TurnInput::UserInput { content, client_id } = input else {
continue;
};
self.emit_user_input_turn_item(turn_context, content, client_id.clone())
.await;
emitted = true;
}
if emitted {
self.ensure_rollout_materialized().await;
}
}
async fn emit_user_input_turn_item(
&self,
turn_context: &TurnContext,
input: &[UserInput],
client_id: Option<String>,
) {
let mut user_message_item = UserMessageItem::new(input);
user_message_item.client_id = client_id;
let turn_item = TurnItem::UserMessage(user_message_item);
self.emit_turn_item_started(turn_context, &turn_item).await;
self.emit_turn_item_completed(turn_context, turn_item).await;
self.ensure_rollout_materialized().await;
}
pub(crate) async fn notify_stream_error(
codex-rs/core/src/session/tests.rs
+131 -32
View File
@@ -16,6 +16,9 @@ use crate::skills::SkillRenderSideEffects;
use crate::skills::render::SkillMetadataBudget;
use crate::test_support::models_manager_with_provider;
use crate::tools::format_exec_output_str;
use codex_app_server_protocol::ThreadItem as AppServerThreadItem;
use codex_app_server_protocol::UserInput as AppServerUserInput;
use codex_app_server_protocol::build_turns_from_rollout_items;
use codex_config::ConfigLayerStack;
use codex_config::ConfigLayerStackOrdering;
use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID;
@@ -146,7 +149,6 @@ use codex_protocol::protocol::ThreadSettingsOverrides;
use codex_protocol::protocol::TokenCountEvent;
use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TokenUsageInfo;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnCompleteEvent;
use codex_protocol::protocol::TurnStartedEvent;
use codex_protocol::protocol::UserMessageEvent;
@@ -439,13 +441,25 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled()
}
#[tokio::test]
async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
async fn interrupting_regular_turn_waiting_on_startup_prewarm_persists_user_item() {
let (mut sess, tc, rx) = make_session_and_context_with_rx().await;
let rollout_path = attach_thread_persistence(
Arc::get_mut(&mut sess).expect("session should not have additional references"),
)
.await;
let (_tx, startup_prewarm_rx) = tokio::sync::oneshot::channel::<()>();
let handle = tokio::spawn(async move {
let _ = startup_prewarm_rx.await;
Ok(test_model_client_session())
});
let prompt = "persist this interrupted prompt";
let input = vec![TurnInput::UserInput {
content: vec![UserInput::Text {
text: prompt.to_string(),
text_elements: Vec::new(),
}],
client_id: None,
}];
sess.set_session_startup_prewarm(
crate::session_startup_prewarm::SessionStartupPrewarmHandle::new(
@@ -455,12 +469,8 @@ async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted
),
)
.await;
sess.spawn_task(
Arc::clone(&tc),
Vec::new(),
crate::tasks::RegularTask::new(),
)
.await;
sess.spawn_task(Arc::clone(&tc), input, crate::tasks::RegularTask::new())
.await;
let first = tokio::time::timeout(std::time::Duration::from_millis(200), rx.recv())
.await
@@ -471,31 +481,120 @@ async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted
EventMsg::TurnStarted(TurnStartedEvent { turn_id, .. }) if turn_id == tc.sub_id
));
sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
tokio::time::timeout(
std::time::Duration::from_secs(1),
sess.abort_all_tasks(TurnAbortReason::Interrupted),
)
.await
.expect("interrupt should not wait for startup hooks");
let marker_evt = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("expected turn aborted marker event")
.expect("channel open");
assert!(matches!(marker_evt.msg, EventMsg::RawResponseItem(_)));
let second = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("expected turn aborted event")
.expect("channel open");
let EventMsg::TurnAborted(TurnAbortedEvent {
turn_id,
reason,
completed_at,
duration_ms,
}) = second.msg
else {
panic!("expected turn aborted event");
let turn_aborted = loop {
let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("expected turn aborted event")
.expect("channel open");
if let EventMsg::TurnAborted(event) = event.msg {
break event;
}
};
assert_eq!(turn_id, Some(tc.sub_id.clone()));
assert_eq!(reason, TurnAbortReason::Interrupted);
assert!(completed_at.is_some());
assert!(duration_ms.is_some());
assert_eq!(turn_aborted.turn_id, Some(tc.sub_id.clone()));
assert_eq!(turn_aborted.reason, TurnAbortReason::Interrupted);
assert!(turn_aborted.completed_at.is_some());
assert!(turn_aborted.duration_ms.is_some());
let is_turn_aborted_marker = |item: &ResponseItem| {
let ResponseItem::Message { content, .. } = item else {
return false;
};
content.iter().any(|content_item| {
let ContentItem::InputText { text } = content_item else {
return false;
};
TurnAborted::matches_text(text)
})
};
let live_history = strip_metadata_from_items(sess.clone_history().await.raw_items());
assert!(!live_history.contains(&user_message(prompt)));
assert_eq!(
live_history
.iter()
.filter(|item| is_turn_aborted_marker(item))
.count(),
1
);
let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path)
.await
.expect("read rollout history")
else {
panic!("expected resumed rollout history");
};
let user_message_positions = resumed
.history
.iter()
.enumerate()
.filter_map(|(index, item)| {
matches!(
item,
RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == prompt
)
.then_some(index)
})
.collect::<Vec<_>>();
let turn_aborted_positions = resumed
.history
.iter()
.enumerate()
.filter_map(|(index, item)| {
matches!(
item,
RolloutItem::EventMsg(EventMsg::TurnAborted(event))
if event.turn_id.as_deref() == Some(tc.sub_id.as_str())
&& event.reason == TurnAbortReason::Interrupted
)
.then_some(index)
})
.collect::<Vec<_>>();
assert_eq!(user_message_positions.len(), 1);
assert_eq!(turn_aborted_positions.len(), 1);
assert!(user_message_positions[0] < turn_aborted_positions[0]);
assert!(!resumed.history.iter().any(
|item| matches!(item, RolloutItem::ResponseItem(item) if item == &user_message(prompt))
));
let turns = build_turns_from_rollout_items(&resumed.history);
let interrupted_turn = turns
.iter()
.find(|turn| turn.id == tc.sub_id)
.expect("interrupted turn should be reconstructed");
let transcript_prompts = interrupted_turn
.items
.iter()
.filter_map(|item| match item {
AppServerThreadItem::UserMessage { content, .. } => Some(content),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(transcript_prompts.len(), 1);
assert!(matches!(
transcript_prompts[0].as_slice(),
[AppServerUserInput::Text { text, .. }] if text == prompt
));
let (resumed_session, _resumed_turn_context) = make_session_and_context().await;
resumed_session
.record_initial_history(InitialHistory::Resumed(resumed))
.await;
let resumed_history =
strip_metadata_from_items(resumed_session.clone_history().await.raw_items());
assert!(!resumed_history.contains(&user_message(prompt)));
assert_eq!(
resumed_history
.iter()
.filter(|item| is_turn_aborted_marker(item))
.count(),
1
);
}
fn test_model_client_session() -> crate::client::ModelClientSession {
codex-rs/core/src/session/turn.rs
+88 -2
View File
@@ -18,6 +18,7 @@ use crate::compact_remote_v2::run_inline_remote_auto_compact_task as run_inline_
use crate::connectors;
use crate::context::ContextualUserFragment;
use crate::feedback_tags;
use crate::hook_runtime::UserInputEventDisposition;
use crate::hook_runtime::inspect_pending_input;
use crate::hook_runtime::record_additional_contexts;
use crate::hook_runtime::record_pending_input;
@@ -53,6 +54,7 @@ use crate::stream_events_utils::last_assistant_message_from_item;
use crate::stream_events_utils::mark_thread_memory_mode_polluted_if_external_context;
use crate::stream_events_utils::raw_assistant_output_text_from_item;
use crate::stream_events_utils::record_completed_response_item_with_finalized_facts;
use crate::tasks::InitialUserInputEvents;
use crate::tasks::emit_compact_metric;
use crate::tools::ToolRouter;
use crate::tools::context::SharedTurnDiffTracker;
@@ -146,6 +148,7 @@ pub(crate) async fn run_turn(
input: Vec<TurnInput>,
prewarmed_client_session: Option<ModelClientSession>,
cancellation_token: CancellationToken,
initial_user_input_events: Option<&InitialUserInputEvents>,
) -> CodexResult<Option<String>> {
let mut client_session =
prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session());
@@ -184,10 +187,25 @@ pub(crate) async fn run_turn(
};
if run_pending_session_start_hooks(&sess, &turn_context).await {
if let Some(initial_user_input_events) = initial_user_input_events {
initial_user_input_events.set_abort_input(Vec::new());
}
return Ok(None);
}
let mut can_drain_pending_input = input.is_empty();
if run_hooks_and_record_inputs(&sess, &turn_context, &input).await {
let initial_input_stopped = if let Some(initial_user_input_events) = initial_user_input_events {
run_hooks_and_record_initial_inputs(&sess, &turn_context, &input, initial_user_input_events)
.await
} else {
run_hooks_and_record_inputs(
&sess,
&turn_context,
&input,
UserInputEventDisposition::Emit,
)
.await
};
if initial_input_stopped {
return Ok(None);
}
@@ -232,7 +250,14 @@ pub(crate) async fn run_turn(
Vec::new()
};
if run_hooks_and_record_inputs(&sess, &turn_context, &pending_input).await {
if run_hooks_and_record_inputs(
&sess,
&turn_context,
&pending_input,
UserInputEventDisposition::Emit,
)
.await
{
break;
}
@@ -483,6 +508,7 @@ async fn run_hooks_and_record_inputs(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
input: &[TurnInput],
user_input_event_disposition: UserInputEventDisposition,
) -> bool {
let mut blocked_input = false;
let mut accepted_user_input = false;
@@ -500,6 +526,7 @@ async fn run_hooks_and_record_inputs(
turn_context,
input_item.clone(),
hook_outcome.additional_contexts,
user_input_event_disposition,
)
.await;
}
@@ -507,6 +534,65 @@ async fn run_hooks_and_record_inputs(
blocked_input && !accepted_user_input
}
#[instrument(level = "trace", skip_all)]
async fn run_hooks_and_record_initial_inputs(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
input: &[TurnInput],
initial_user_input_events: &InitialUserInputEvents,
) -> bool {
let mut blocked_input = false;
let mut accepted_user_input = false;
let mut accepted: Vec<(TurnInput, Vec<String>)> = Vec::new();
for (index, input_item) in input.iter().enumerate() {
let hook_outcome = inspect_pending_input(sess, turn_context, input_item).await;
if hook_outcome.should_stop {
blocked_input = true;
let abort_input = accepted
.iter()
.map(|(input, _additional_contexts)| input.clone())
.chain(input.iter().skip(index + 1).cloned())
.collect();
initial_user_input_events.set_abort_input(abort_input);
record_additional_contexts(sess, turn_context, hook_outcome.additional_contexts).await;
} else {
if matches!(input_item, TurnInput::UserInput { content, .. } if !content.is_empty()) {
accepted_user_input = true;
}
accepted.push((input_item.clone(), hook_outcome.additional_contexts));
let abort_input = accepted
.iter()
.map(|(input, _additional_contexts)| input.clone())
.chain(input.iter().skip(index + 1).cloned())
.collect();
initial_user_input_events.set_abort_input(abort_input);
}
}
let accepted_input = accepted
.iter()
.map(|(input, _additional_contexts)| input.clone())
.collect();
if !initial_user_input_events
.record(Arc::clone(sess), Arc::clone(turn_context), accepted_input)
.await
{
return true;
}
for (input, additional_contexts) in accepted {
record_pending_input(
sess,
turn_context,
input,
additional_contexts,
UserInputEventDisposition::AlreadyEmitted,
)
.await;
}
blocked_input && !accepted_user_input
}
#[instrument(level = "trace", skip_all)]
async fn build_skills_and_plugins(
sess: &Arc<Session>,
codex-rs/core/src/state/turn.rs
+1
View File
@@ -73,6 +73,7 @@ pub(crate) struct RunningTask {
pub(crate) done: Arc<Notify>,
pub(crate) kind: TaskKind,
pub(crate) task: Arc<dyn AnySessionTask>,
pub(crate) abort_input: Vec<crate::session::TurnInput>,
pub(crate) cancellation_token: CancellationToken,
pub(crate) handle: AbortOnDropHandle<()>,
pub(crate) turn_context: Arc<TurnContext>,
codex-rs/core/src/tasks/mod.rs
+35 -8
View File
@@ -25,6 +25,7 @@ use tracing::warn;
use crate::codex_thread::BackgroundTerminalInfo;
use crate::config::Config;
use crate::context::ContextualUserFragment;
use crate::hook_runtime::UserInputEventDisposition;
use crate::hook_runtime::inspect_pending_input;
use crate::hook_runtime::record_additional_contexts;
use crate::hook_runtime::record_pending_input;
@@ -58,6 +59,7 @@ use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::models::ContentItem;
pub(crate) use compact::CompactTask;
pub(crate) use regular::InitialUserInputEvents;
pub(crate) use regular::RegularTask;
pub(crate) use review::ReviewTask;
pub(crate) use user_shell::UserShellCommandMode;
@@ -241,14 +243,17 @@ pub(crate) trait SessionTask: Send + Sync + 'static {
///
/// The default implementation is a no-op; override this if additional
/// teardown or notifications are required once
/// [`Session::abort_all_tasks`] cancels the task.
fn abort(
&self,
/// [`Session::abort_all_tasks`] cancels the task. `input` is the same initial input passed to
/// [`SessionTask::run`], retained for cleanup that must handle cancellation before the run
/// future is first polled.
fn abort<'a>(
&'a self,
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
) -> impl std::future::Future<Output = ()> + Send {
input: &'a [TurnInput],
) -> impl std::future::Future<Output = ()> + Send + 'a {
async move {
let _ = (session, ctx);
let _ = (session, ctx, input);
}
}
}
@@ -270,6 +275,7 @@ pub(crate) trait AnySessionTask: Send + Sync + 'static {
&'a self,
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
input: &'a [TurnInput],
) -> BoxFuture<'a, ()>;
}
@@ -305,8 +311,9 @@ where
&'a self,
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
input: &'a [TurnInput],
) -> BoxFuture<'a, ()> {
Box::pin(SessionTask::abort(self, session, ctx))
Box::pin(SessionTask::abort(self, session, ctx, input))
}
}
@@ -379,6 +386,7 @@ impl Session {
));
let ctx = Arc::clone(&turn_context);
let task_for_run = Arc::clone(&task);
let abort_input = input.clone();
let task_input = input;
let task_cancellation_token = cancellation_token.child_token();
// Task-owned turn spans keep a core-owned span open for the
@@ -441,6 +449,7 @@ impl Session {
handle: AbortOnDropHandle::new(handle),
kind: task_kind,
task,
abort_input,
cancellation_token,
turn_context: Arc::clone(&turn_context),
turn_extension_data,
@@ -617,6 +626,7 @@ impl Session {
&turn_context,
pending_input_item,
hook_outcome.additional_contexts,
UserInputEventDisposition::Emit,
)
.await;
}
@@ -847,14 +857,31 @@ impl Session {
}
}
task.handle.abort();
let mut handle = task.handle;
handle.abort();
// Wait for cancellation to finish before task-specific abort cleanup reads state that the
// run future may update between cancellable awaits. Keep this bounded because Tokio cannot
// force a task to stop while it is executing non-yielding code.
if tokio::time::timeout(
Duration::from_millis(GRACEFULL_INTERRUPTION_TIMEOUT_MS),
&mut handle,
)
.await
.is_err()
{
warn!("task {sub_id} did not stop after forced abort");
}
let session_ctx = Arc::new(SessionTaskContext::new(
Arc::clone(self),
Arc::clone(&task.turn_extension_data),
));
session_task
.abort(session_ctx, Arc::clone(&task.turn_context))
.abort(
session_ctx,
Arc::clone(&task.turn_context),
&task.abort_input,
)
.await;
if reason == TurnAbortReason::Interrupted
codex-rs/core/src/tasks/regular.rs
+150 -13
View File
@@ -1,5 +1,9 @@
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use crate::session::TurnInput;
@@ -16,12 +20,136 @@ use super::SessionTask;
use super::SessionTaskContext;
use super::SessionTaskResult;
pub(crate) struct InitialUserInputEvents {
abort_input: Mutex<Option<Vec<TurnInput>>>,
turn_start_started: AtomicBool,
turn_start_recorded: watch::Sender<bool>,
started: AtomicBool,
recorded: watch::Sender<bool>,
}
impl Default for InitialUserInputEvents {
fn default() -> Self {
let (recorded, _receiver) = watch::channel(false);
let (turn_start_recorded, _receiver) = watch::channel(false);
Self {
abort_input: Mutex::new(None),
turn_start_started: AtomicBool::new(false),
turn_start_recorded,
started: AtomicBool::new(false),
recorded,
}
}
}
impl InitialUserInputEvents {
/// Emit the regular-turn boundary exactly once, even when cancellation transfers ownership
/// from the task future to abort cleanup.
async fn record_turn_start(
&self,
session: Arc<crate::session::session::Session>,
ctx: Arc<TurnContext>,
) {
let mut recorded = self.turn_start_recorded.subscribe();
let claimed = self
.turn_start_started
.compare_exchange(
/*current*/ false,
/*new*/ true,
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok();
if claimed {
let recorded = self.turn_start_recorded.clone();
tokio::spawn(async move {
session
.send_event(
ctx.as_ref(),
EventMsg::TurnStarted(TurnStartedEvent {
turn_id: ctx.sub_id.clone(),
trace_id: ctx.trace_id.clone(),
started_at: ctx.turn_timing_state.started_at_unix_secs().await,
model_context_window: ctx.model_context_window(),
collaboration_mode_kind: ctx.collaboration_mode.mode,
}),
)
.await;
session
.set_server_reasoning_included(/*included*/ false)
.await;
recorded.send_replace(true);
});
}
while !*recorded.borrow_and_update() {
// `self.turn_start_recorded` keeps the channel open for this task's lifetime.
let _ = recorded.changed().await;
}
}
/// Publish the hook-filtered input that an interrupt should preserve if it wins before normal
/// persistence. This is synchronous so a completed hook decision is visible before the next
/// cancellable await.
pub(crate) fn set_abort_input(&self, input: Vec<TurnInput>) {
*self
.abort_input
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(input);
}
fn abort_input(&self, original_input: &[TurnInput]) -> Vec<TurnInput> {
self.abort_input
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
.unwrap_or_else(|| original_input.to_vec())
}
/// Record one durable copy of the initial user items and report whether this caller won the
/// right to choose which items were recorded.
pub(crate) async fn record(
&self,
session: Arc<crate::session::session::Session>,
ctx: Arc<TurnContext>,
input: Vec<TurnInput>,
) -> bool {
let mut recorded = self.recorded.subscribe();
let claimed = self
.started
.compare_exchange(
/*current*/ false,
/*new*/ true,
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok();
if claimed {
let recorded = self.recorded.clone();
tokio::spawn(async move {
session
.record_initial_user_input_events(ctx.as_ref(), &input)
.await;
recorded.send_replace(true);
});
}
while !*recorded.borrow_and_update() {
// `self.recorded` keeps the channel open for the duration of this borrow.
let _ = recorded.changed().await;
}
claimed
}
}
#[derive(Default)]
pub(crate) struct RegularTask;
pub(crate) struct RegularTask {
initial_user_input_events: InitialUserInputEvents,
}
impl RegularTask {
pub(crate) fn new() -> Self {
Self
Self::default()
}
}
@@ -44,18 +172,10 @@ impl SessionTask for RegularTask {
let sess = session.clone_session();
let turn_extension_data = session.turn_extension_data();
let run_turn_span = trace_span!("run_turn");
// Regular turns emit `TurnStarted` inline so first-turn lifecycle does
// not wait on startup prewarm resolution.
self.initial_user_input_events
.record_turn_start(Arc::clone(&sess), Arc::clone(&ctx))
.await;
let prewarmed_client_session = async {
let event = EventMsg::TurnStarted(TurnStartedEvent {
turn_id: ctx.sub_id.clone(),
trace_id: ctx.trace_id.clone(),
started_at: ctx.turn_timing_state.started_at_unix_secs().await,
model_context_window: ctx.model_context_window(),
collaboration_mode_kind: ctx.collaboration_mode.mode,
});
sess.send_event(ctx.as_ref(), event).await;
sess.set_server_reasoning_included(/*included*/ false).await;
sess.consume_startup_prewarm_for_regular_turn(&cancellation_token)
.await
}
@@ -70,6 +190,7 @@ impl SessionTask for RegularTask {
};
let mut next_input = input;
let mut prewarmed_client_session = prewarmed_client_session;
let mut initial_user_input_events = Some(&self.initial_user_input_events);
loop {
let last_agent_message = run_turn(
Arc::clone(&sess),
@@ -78,6 +199,7 @@ impl SessionTask for RegularTask {
next_input,
prewarmed_client_session.take(),
cancellation_token.child_token(),
initial_user_input_events.take(),
)
.instrument(run_turn_span.clone())
.await?;
@@ -87,4 +209,19 @@ impl SessionTask for RegularTask {
next_input = Vec::new();
}
}
async fn abort(
&self,
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
input: &[TurnInput],
) {
self.initial_user_input_events
.record_turn_start(session.clone_session(), Arc::clone(&ctx))
.await;
let input = self.initial_user_input_events.abort_input(input);
self.initial_user_input_events
.record(session.clone_session(), ctx, input)
.await;
}
}
codex-rs/core/src/tasks/review.rs
+6 -1
View File
@@ -88,7 +88,12 @@ impl SessionTask for ReviewTask {
Ok(None)
}
async fn abort(&self, session: Arc<SessionTaskContext>, ctx: Arc<TurnContext>) {
async fn abort(
&self,
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
_input: &[TurnInput],
) {
exit_review_mode(session.clone_session(), /*review_output*/ None, ctx).await;
}
}
codex-rs/core/tests/suite/hooks.rs
+119
View File
@@ -22,6 +22,7 @@ use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::RolloutLine;
use codex_protocol::user_input::UserInput;
use codex_utils_absolute_path::AbsolutePathBuf;
use core_test_support::fs_wait;
use core_test_support::hooks::trust_discovered_hooks;
use core_test_support::hooks::trust_hooks;
use core_test_support::managed_network_requirements_loader;
@@ -259,6 +260,43 @@ if payload.get("prompt") == {blocked_prompt_json}:
Ok(())
}
fn write_blocking_user_prompt_submit_hook(home: &Path) -> Result<()> {
let script_path = home.join("blocking_user_prompt_submit_hook.py");
let started_path = home.join("blocking_user_prompt_submit_hook.started");
let release_path = home.join("blocking_user_prompt_submit_hook.release");
let script = format!(
r#"from pathlib import Path
import json
import sys
import time
json.load(sys.stdin)
Path(r"{started_path}").touch()
release_path = Path(r"{release_path}")
while not release_path.exists():
time.sleep(0.01)
print(json.dumps({{}}))
"#,
started_path = started_path.display(),
release_path = release_path.display(),
);
let hooks = serde_json::json!({
"hooks": {
"UserPromptSubmit": [{
"hooks": [{
"type": "command",
"command": format!("python3 {}", script_path.display()),
"statusMessage": "waiting in user prompt submit hook",
}]
}]
}
});
fs::write(script_path, script).context("write blocking user prompt submit hook")?;
fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?;
Ok(())
}
fn write_session_start_and_user_prompt_submit_order_hooks(home: &Path) -> Result<()> {
let session_start_script_path = home.join("session_start_order_hook.py");
let user_prompt_submit_script_path = home.join("user_prompt_submit_order_hook.py");
@@ -921,6 +959,24 @@ fn rollout_hook_prompt_texts(text: &str) -> Result<Vec<String>> {
Ok(texts)
}
fn rollout_user_message_event_texts(text: &str) -> Result<Vec<String>> {
text.lines()
.filter(|line| !line.trim().is_empty())
.filter_map(|line| {
let rollout: Result<RolloutLine> =
serde_json::from_str(line).context("parse rollout line");
match rollout {
Ok(RolloutLine {
item: RolloutItem::EventMsg(EventMsg::UserMessage(event)),
..
}) => Some(Ok(event.message)),
Ok(_) => None,
Err(err) => Some(Err(err)),
}
})
.collect()
}
fn request_hook_prompt_texts(
request: &core_test_support::responses::ResponsesRequest,
) -> Vec<String> {
@@ -1753,6 +1809,14 @@ async fn blocked_user_prompt_submit_persists_additional_context_for_next_turn()
"second request should include the accepted prompt",
);
let rollout_path = test.codex.rollout_path().expect("rollout path");
let rollout_text = fs::read_to_string(rollout_path)?;
assert_eq!(
rollout_user_message_event_texts(&rollout_text)?,
vec!["second prompt".to_string()],
"blocked prompts should not become transcript user items",
);
let hook_inputs = read_user_prompt_submit_hook_inputs(test.codex_home_path())?;
assert_eq!(hook_inputs.len(), 2);
assert_eq!(
@@ -1780,6 +1844,61 @@ async fn blocked_user_prompt_submit_persists_additional_context_for_next_turn()
Ok(())
}
#[tokio::test]
async fn interrupt_during_user_prompt_submit_hook_persists_one_transcript_item() -> Result<()> {
skip_if_host_windows!(Ok(()));
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let mut builder = test_codex()
.with_pre_build_hook(|home| {
write_blocking_user_prompt_submit_hook(home)
.expect("failed to write blocking user prompt submit hook fixture");
})
.with_config(trust_discovered_hooks);
let test = builder.build(&server).await?;
let prompt = "persist the prompt interrupted during its hook";
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: prompt.to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
})
.await?;
fs_wait::wait_for_path_exists(
test.codex_home_path()
.join("blocking_user_prompt_submit_hook.started"),
Duration::from_secs(5),
)
.await?;
test.codex.submit(Op::Interrupt).await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnAborted(_))
})
.await;
fs::write(
test.codex_home_path()
.join("blocking_user_prompt_submit_hook.release"),
b"release",
)?;
let rollout_path = test.codex.rollout_path().expect("rollout path");
let rollout_text = fs::read_to_string(rollout_path)?;
assert_eq!(
rollout_user_message_event_texts(&rollout_text)?,
vec![prompt.to_string()],
);
Ok(())
}
#[tokio::test]
async fn blocked_queued_prompt_does_not_strand_earlier_accepted_prompt() -> Result<()> {
skip_if_no_network!(Ok(()));