fix: emit thread stop lifecycle on implicit shutdown (#22482)

## Why

The thread lifecycle contributor hooks from #22476 should observe every
session teardown. The explicit `Op::Shutdown` path already emitted
`on_thread_stop`, but when `submission_loop` exited because its
submission channel closed, it only tore down runtime services. That
meant extensions could miss the thread-stop lifecycle signal on implicit
runtime shutdown.

## What Changed

- Split shared runtime teardown into `shutdown_runtime_services(...)`.
- Split thread-stop lifecycle emission into
`emit_thread_stop_lifecycle(...)`.
- Reused those helpers from both explicit shutdown and the channel-close
shutdown path.
- Tracked whether `Op::Shutdown` was received so the explicit path does
not double-emit lifecycle events after it exits the loop.
- Added a regression test that closes the submission channel and asserts
`ThreadLifecycleContributor::on_thread_stop` runs once with the expected
thread/session stores.

## Testing

- `cargo test -p codex-core
submission_loop_channel_close_emits_thread_stop_lifecycle`
This commit is contained in:
jif-oai
2026-05-13 14:19:57 +02:00
committed by GitHub
Unverified
parent 27e67a8c2a
commit fcc2a92743
2 changed files with 133 additions and 21 deletions
+23 -21
View File
@@ -617,7 +617,7 @@ pub async fn set_thread_memory_mode(sess: &Arc<Session>, sub_id: String, mode: T
}
}
pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
async fn shutdown_session_runtime(sess: &Arc<Session>) {
sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
let _ = sess.conversation.shutdown().await;
sess.services
@@ -630,6 +630,20 @@ pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
};
mcp_shutdown.await;
sess.guardian_review_session.shutdown().await;
}
fn emit_thread_stop_lifecycle(sess: &Session) {
for contributor in sess.services.extensions.thread_lifecycle_contributors() {
contributor.on_thread_stop(codex_extension_api::ThreadStopInput {
thread_id: sess.conversation_id,
session_store: &sess.services.session_extension_data,
thread_store: &sess.services.thread_extension_data,
});
}
}
pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
shutdown_session_runtime(sess).await;
info!("Shutting down Codex instance");
let history = sess.clone_history().await;
let turn_count = history
@@ -643,13 +657,7 @@ pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
&[],
);
for contributor in sess.services.extensions.thread_lifecycle_contributors() {
contributor.on_thread_stop(codex_extension_api::ThreadStopInput {
thread_id: sess.conversation_id,
session_store: &sess.services.session_extension_data,
thread_store: &sess.services.thread_extension_data,
});
}
emit_thread_stop_lifecycle(sess.as_ref());
// Gracefully flush and shutdown thread persistence on session end so tests
// that inspect durable state do not race with the background writer.
@@ -722,6 +730,7 @@ pub(super) async fn submission_loop(
rx_sub: Receiver<Submission>,
) {
// To break out of this loop, send Op::Shutdown.
let mut shutdown_received = false;
while let Ok(sub) = rx_sub.recv().await {
debug!(?sub, "Submission");
let dispatch_span = submission_dispatch_span(&sub);
@@ -894,23 +903,16 @@ pub(super) async fn submission_loop(
.instrument(dispatch_span)
.await;
if should_exit {
shutdown_received = true;
break;
}
}
// If the submission loop exits because the channel closed without an
// explicit shutdown op, still run process teardown for child processes
// owned by this session.
sess.services
.unified_exec_manager
.terminate_all_processes()
.await;
let mcp_shutdown = {
let mut manager = sess.services.mcp_connection_manager.write().await;
manager.begin_shutdown()
};
mcp_shutdown.await;
// Also drain cached guardian state on this implicit shutdown path.
sess.guardian_review_session.shutdown().await;
// explicit shutdown op, still run session teardown.
if !shutdown_received {
shutdown_session_runtime(&sess).await;
emit_thread_stop_lifecycle(sess.as_ref());
}
debug!("Agent loop exited");
}
+110
View File
@@ -5196,6 +5196,116 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() {
);
}
#[tokio::test]
async fn submission_loop_channel_close_emits_thread_stop_lifecycle() {
struct SessionStopMarker;
struct ThreadStopMarker;
struct ThreadStopRecorder {
calls: Arc<std::sync::atomic::AtomicUsize>,
expected_thread_id: ThreadId,
}
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for ThreadStopRecorder {
fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
assert_eq!(self.expected_thread_id, input.thread_id);
assert!(input.session_store.get::<SessionStopMarker>().is_some());
assert!(input.thread_store.get::<ThreadStopMarker>().is_some());
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
}
let (mut session, turn_context) = make_session_and_context().await;
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let mut builder = codex_extension_api::ExtensionRegistryBuilder::<crate::config::Config>::new();
builder.thread_lifecycle_contributor(Arc::new(ThreadStopRecorder {
calls: Arc::clone(&calls),
expected_thread_id: session.conversation_id,
}));
session.services.extensions = Arc::new(builder.build());
session
.services
.session_extension_data
.insert(SessionStopMarker);
session
.services
.thread_extension_data
.insert(ThreadStopMarker);
let (tx_sub, rx_sub) = async_channel::bounded(1);
drop(tx_sub);
let session = Arc::new(session);
submission_loop(session, Arc::clone(&turn_context.config), rx_sub).await;
assert_eq!(1, calls.load(std::sync::atomic::Ordering::SeqCst));
}
#[tokio::test]
async fn submission_loop_channel_close_aborts_active_turn_before_thread_stop_lifecycle() {
struct LifecycleRecorder {
calls: Arc<std::sync::Mutex<Vec<&'static str>>>,
expected_thread_id: ThreadId,
expected_turn_id: String,
}
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for LifecycleRecorder {
fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
assert_eq!(self.expected_thread_id, input.thread_id);
self.calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push("thread_stop");
}
}
impl codex_extension_api::TurnLifecycleContributor for LifecycleRecorder {
fn on_turn_abort(&self, input: codex_extension_api::TurnAbortInput<'_>) {
assert_eq!(self.expected_thread_id, input.thread_id);
assert_eq!(self.expected_turn_id, input.turn_id);
assert_eq!(TurnAbortReason::Interrupted, input.reason);
self.calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push("turn_abort");
}
}
let (mut session, turn_context) = make_session_and_context().await;
let calls = Arc::new(std::sync::Mutex::new(Vec::new()));
let recorder = Arc::new(LifecycleRecorder {
calls: Arc::clone(&calls),
expected_thread_id: session.conversation_id,
expected_turn_id: turn_context.sub_id.clone(),
});
let mut builder = codex_extension_api::ExtensionRegistryBuilder::<crate::config::Config>::new();
builder.thread_lifecycle_contributor(recorder.clone());
builder.turn_lifecycle_contributor(recorder);
session.services.extensions = Arc::new(builder.build());
let session = Arc::new(session);
session
.spawn_task(
Arc::new(turn_context),
Vec::new(),
NeverEndingTask {
kind: TaskKind::Regular,
listen_to_cancellation_token: true,
},
)
.await;
let (tx_sub, rx_sub) = async_channel::bounded(1);
drop(tx_sub);
submission_loop(Arc::clone(&session), session.get_config().await, rx_sub).await;
assert_eq!(
vec!["turn_abort", "thread_stop"],
*calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
);
}
#[tokio::test]
async fn shutdown_and_wait_allows_multiple_waiters() {
let (session, _turn_context) = make_session_and_context().await;