mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Make extension lifecycle hooks async (#23291)
## Why Extension lifecycle hooks sit on the host/extension boundary, but the current trait surface only allows synchronous callbacks. That forces extensions that need to seed, rehydrate, observe, or flush extension-owned state during thread and turn transitions to either block inside the callback or move async work into separate host plumbing. This PR makes those lifecycle callbacks awaitable so extension implementations can perform async work directly at the lifecycle point where the host already has the relevant session, thread, or turn stores available. ## What changed - Makes `ThreadLifecycleContributor` and `TurnLifecycleContributor` async in `codex-extension-api`. - Awaits thread start/resume/stop and turn start/stop/abort lifecycle callbacks from `codex-core`. - Updates the guardian and memories extensions to implement the async lifecycle trait surface. - Updates the existing lifecycle tests to use async contributor implementations. - Adds `async-trait` to the crates that now expose or implement these async object-safe lifecycle traits. ## Testing - Existing `codex-core` lifecycle tests were updated to cover async implementations for thread stop and turn abort ordering.
This commit is contained in:
committed by
GitHub
Unverified
parent
a80f07ec4a
commit
9531e932ef
@@ -144,7 +144,7 @@ impl CodexThread {
|
||||
self.codex.session_loop_termination.clone().await;
|
||||
}
|
||||
|
||||
pub(crate) fn emit_thread_resume_lifecycle(&self) {
|
||||
pub(crate) async fn emit_thread_resume_lifecycle(&self) {
|
||||
for contributor in self
|
||||
.codex
|
||||
.session
|
||||
@@ -152,10 +152,12 @@ impl CodexThread {
|
||||
.extensions
|
||||
.thread_lifecycle_contributors()
|
||||
{
|
||||
contributor.on_thread_resume(codex_extension_api::ThreadResumeInput {
|
||||
session_store: &self.codex.session.services.session_extension_data,
|
||||
thread_store: &self.codex.session.services.thread_extension_data,
|
||||
});
|
||||
contributor
|
||||
.on_thread_resume(codex_extension_api::ThreadResumeInput {
|
||||
session_store: &self.codex.session.services.session_extension_data,
|
||||
thread_store: &self.codex.session.services.thread_extension_data,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -638,12 +638,14 @@ async fn shutdown_session_runtime(sess: &Arc<Session>) {
|
||||
sess.guardian_review_session.shutdown().await;
|
||||
}
|
||||
|
||||
fn emit_thread_stop_lifecycle(sess: &Session) {
|
||||
async fn emit_thread_stop_lifecycle(sess: &Session) {
|
||||
for contributor in sess.services.extensions.thread_lifecycle_contributors() {
|
||||
contributor.on_thread_stop(codex_extension_api::ThreadStopInput {
|
||||
session_store: &sess.services.session_extension_data,
|
||||
thread_store: &sess.services.thread_extension_data,
|
||||
});
|
||||
contributor
|
||||
.on_thread_stop(codex_extension_api::ThreadStopInput {
|
||||
session_store: &sess.services.session_extension_data,
|
||||
thread_store: &sess.services.thread_extension_data,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -662,7 +664,7 @@ pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
|
||||
&[],
|
||||
);
|
||||
|
||||
emit_thread_stop_lifecycle(sess.as_ref());
|
||||
emit_thread_stop_lifecycle(sess.as_ref()).await;
|
||||
|
||||
// Gracefully flush and shutdown thread persistence on session end so tests
|
||||
// that inspect durable state do not race with the background writer.
|
||||
@@ -917,7 +919,7 @@ pub(super) async fn submission_loop(
|
||||
// explicit shutdown op, still run session teardown.
|
||||
if !shutdown_received {
|
||||
shutdown_session_runtime(&sess).await;
|
||||
emit_thread_stop_lifecycle(sess.as_ref());
|
||||
emit_thread_stop_lifecycle(sess.as_ref()).await;
|
||||
}
|
||||
debug!("Agent loop exited");
|
||||
}
|
||||
|
||||
@@ -877,7 +877,7 @@ impl Session {
|
||||
config: config.as_ref(),
|
||||
session_store: &session_extension_data,
|
||||
thread_store: &thread_extension_data,
|
||||
});
|
||||
}).await;
|
||||
}
|
||||
|
||||
let services = SessionServices {
|
||||
|
||||
@@ -5683,8 +5683,9 @@ async fn submission_loop_channel_close_emits_thread_stop_lifecycle() {
|
||||
expected_thread_id: ThreadId,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for ThreadStopRecorder {
|
||||
fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
|
||||
async fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
|
||||
assert_eq!(
|
||||
self.expected_thread_id.to_string(),
|
||||
input.thread_store.level_id()
|
||||
@@ -5728,8 +5729,9 @@ async fn submission_loop_channel_close_aborts_active_turn_before_thread_stop_lif
|
||||
expected_turn_id: String,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for LifecycleRecorder {
|
||||
fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
|
||||
async fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
|
||||
assert_eq!(
|
||||
self.expected_thread_id.to_string(),
|
||||
input.thread_store.level_id()
|
||||
@@ -5741,8 +5743,9 @@ async fn submission_loop_channel_close_aborts_active_turn_before_thread_stop_lif
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl codex_extension_api::TurnLifecycleContributor for LifecycleRecorder {
|
||||
fn on_turn_abort(&self, input: codex_extension_api::TurnAbortInput<'_>) {
|
||||
async fn on_turn_abort(&self, input: codex_extension_api::TurnAbortInput<'_>) {
|
||||
assert_eq!(
|
||||
self.expected_thread_id.to_string(),
|
||||
input.thread_store.level_id()
|
||||
|
||||
@@ -4,38 +4,44 @@ use codex_protocol::protocol::TurnAbortReason;
|
||||
use crate::session::session::Session;
|
||||
|
||||
impl Session {
|
||||
pub(super) fn emit_turn_start_lifecycle(&self, turn_store: &ExtensionData) {
|
||||
pub(super) async fn emit_turn_start_lifecycle(&self, turn_store: &ExtensionData) {
|
||||
for contributor in self.services.extensions.turn_lifecycle_contributors() {
|
||||
contributor.on_turn_start(codex_extension_api::TurnStartInput {
|
||||
session_store: &self.services.session_extension_data,
|
||||
thread_store: &self.services.thread_extension_data,
|
||||
turn_store,
|
||||
});
|
||||
contributor
|
||||
.on_turn_start(codex_extension_api::TurnStartInput {
|
||||
session_store: &self.services.session_extension_data,
|
||||
thread_store: &self.services.thread_extension_data,
|
||||
turn_store,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn emit_turn_stop_lifecycle(&self, turn_store: &ExtensionData) {
|
||||
pub(super) async fn emit_turn_stop_lifecycle(&self, turn_store: &ExtensionData) {
|
||||
for contributor in self.services.extensions.turn_lifecycle_contributors() {
|
||||
contributor.on_turn_stop(codex_extension_api::TurnStopInput {
|
||||
session_store: &self.services.session_extension_data,
|
||||
thread_store: &self.services.thread_extension_data,
|
||||
turn_store,
|
||||
});
|
||||
contributor
|
||||
.on_turn_stop(codex_extension_api::TurnStopInput {
|
||||
session_store: &self.services.session_extension_data,
|
||||
thread_store: &self.services.thread_extension_data,
|
||||
turn_store,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn emit_turn_abort_lifecycle(
|
||||
pub(super) async fn emit_turn_abort_lifecycle(
|
||||
&self,
|
||||
reason: TurnAbortReason,
|
||||
turn_store: &ExtensionData,
|
||||
) {
|
||||
for contributor in self.services.extensions.turn_lifecycle_contributors() {
|
||||
contributor.on_turn_abort(codex_extension_api::TurnAbortInput {
|
||||
reason: reason.clone(),
|
||||
session_store: &self.services.session_extension_data,
|
||||
thread_store: &self.services.thread_extension_data,
|
||||
turn_store,
|
||||
});
|
||||
contributor
|
||||
.on_turn_abort(codex_extension_api::TurnAbortInput {
|
||||
reason: reason.clone(),
|
||||
session_store: &self.services.session_extension_data,
|
||||
thread_store: &self.services.thread_extension_data,
|
||||
turn_store,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -365,7 +365,8 @@ impl Session {
|
||||
turn_state.push_pending_input(item);
|
||||
}
|
||||
}
|
||||
self.emit_turn_start_lifecycle(turn_context.extension_data.as_ref());
|
||||
self.emit_turn_start_lifecycle(turn_context.extension_data.as_ref())
|
||||
.await;
|
||||
|
||||
let turn_extension_data = Arc::clone(&turn_context.extension_data);
|
||||
let mut active = self.active_turn.lock().await;
|
||||
@@ -505,7 +506,8 @@ impl Session {
|
||||
}
|
||||
|
||||
if let Some(turn_context) = turn_context.as_deref() {
|
||||
self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref());
|
||||
self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref())
|
||||
.await;
|
||||
}
|
||||
if (aborted_turn || reason == TurnAbortReason::Interrupted)
|
||||
&& let Err(err) = self
|
||||
@@ -553,7 +555,8 @@ impl Session {
|
||||
self.handle_task_abort(task, reason.clone()).await;
|
||||
}
|
||||
if let Some(turn_context) = turn_context.as_deref() {
|
||||
self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref());
|
||||
self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref())
|
||||
.await;
|
||||
}
|
||||
if let Err(err) = self
|
||||
.goal_runtime_apply(GoalRuntimeEvent::TaskAborted {
|
||||
@@ -756,7 +759,8 @@ impl Session {
|
||||
.time_to_first_token_ms()
|
||||
.await;
|
||||
if should_clear_active_turn {
|
||||
self.emit_turn_stop_lifecycle(turn_context.extension_data.as_ref());
|
||||
self.emit_turn_stop_lifecycle(turn_context.extension_data.as_ref())
|
||||
.await;
|
||||
}
|
||||
if let Err(err) = self
|
||||
.goal_runtime_apply(GoalRuntimeEvent::TurnFinished {
|
||||
|
||||
@@ -1228,7 +1228,7 @@ impl ThreadManagerState {
|
||||
.finalize_thread_spawn(codex, thread_id, tracked_session_source)
|
||||
.await?;
|
||||
if is_resumed_thread {
|
||||
new_thread.thread.emit_thread_resume_lifecycle();
|
||||
new_thread.thread.emit_thread_resume_lifecycle().await;
|
||||
if let Err(err) = new_thread.thread.apply_goal_resume_runtime_effects().await {
|
||||
warn!("failed to apply goal resume runtime effects: {err}");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user