From 9531e932ef33b3834e5dc23fe71ae7ac84e6e213 Mon Sep 17 00:00:00 2001
From: jif-oai <jif@openai.com>
Date: Mon, 18 May 2026 13:53:58 +0200
Subject: [PATCH] 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.
---
 codex-rs/Cargo.lock                           |  2 +
 codex-rs/core/src/codex_thread.rs             | 12 ++---
 codex-rs/core/src/session/handlers.rs         | 16 ++++---
 codex-rs/core/src/session/session.rs          |  2 +-
 codex-rs/core/src/session/tests.rs            |  9 ++--
 codex-rs/core/src/tasks/lifecycle.rs          | 44 +++++++++++--------
 codex-rs/core/src/tasks/mod.rs                | 12 +++--
 codex-rs/core/src/thread_manager.rs           |  2 +-
 codex-rs/ext/extension-api/Cargo.toml         |  1 +
 .../ext/extension-api/src/contributors.rs     | 16 ++++---
 codex-rs/ext/extension-api/src/registry.rs    | 12 ++---
 codex-rs/ext/goal/src/extension.rs            | 10 +++--
 codex-rs/ext/guardian/Cargo.toml              |  1 +
 codex-rs/ext/guardian/src/lib.rs              |  3 +-
 codex-rs/ext/memories/src/extension.rs        |  3 +-
 15 files changed, 86 insertions(+), 59 deletions(-)

diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock
index 850ddbd525..5ced1803fa 100644
--- a/codex-rs/Cargo.lock
+++ b/codex-rs/Cargo.lock
@@ -2829,6 +2829,7 @@ dependencies = [
 name = "codex-extension-api"
 version = "0.0.0"
 dependencies = [
+ "async-trait",
  "codex-protocol",
  "codex-tools",
 ]
@@ -2962,6 +2963,7 @@ dependencies = [
 name = "codex-guardian"
 version = "0.0.0"
 dependencies = [
+ "async-trait",
  "codex-core",
  "codex-extension-api",
  "codex-protocol",
diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs
index 1a2d5ed710..7e9e6e38cd 100644
--- a/codex-rs/core/src/codex_thread.rs
+++ b/codex-rs/core/src/codex_thread.rs
@@ -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;
         }
     }
 
diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs
index a26b9bb3f6..c2ffe2f4e4 100644
--- a/codex-rs/core/src/session/handlers.rs
+++ b/codex-rs/core/src/session/handlers.rs
@@ -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");
 }
diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs
index 3f8276d04d..925ee2df59 100644
--- a/codex-rs/core/src/session/session.rs
+++ b/codex-rs/core/src/session/session.rs
@@ -877,7 +877,7 @@ impl Session {
                     config: config.as_ref(),
                     session_store: &session_extension_data,
                     thread_store: &thread_extension_data,
-                });
+                }).await;
             }
 
             let services = SessionServices {
diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs
index bb858999e4..511dbfeabb 100644
--- a/codex-rs/core/src/session/tests.rs
+++ b/codex-rs/core/src/session/tests.rs
@@ -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()
diff --git a/codex-rs/core/src/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs
index 2fb9017df2..94b51647c9 100644
--- a/codex-rs/core/src/tasks/lifecycle.rs
+++ b/codex-rs/core/src/tasks/lifecycle.rs
@@ -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;
         }
     }
 }
diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs
index 58ead0fd25..7724d5f2d7 100644
--- a/codex-rs/core/src/tasks/mod.rs
+++ b/codex-rs/core/src/tasks/mod.rs
@@ -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 {
diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs
index 7b3d0a9921..e5b4fcf6a2 100644
--- a/codex-rs/core/src/thread_manager.rs
+++ b/codex-rs/core/src/thread_manager.rs
@@ -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}");
             }
diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml
index ceecc48023..15f6e99fa0 100644
--- a/codex-rs/ext/extension-api/Cargo.toml
+++ b/codex-rs/ext/extension-api/Cargo.toml
@@ -14,5 +14,6 @@ doctest = false
 workspace = true
 
 [dependencies]
+async-trait = { workspace = true }
 codex-protocol = { workspace = true }
 codex-tools = { workspace = true }
diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs
index a8c2001200..b03846fd31 100644
--- a/codex-rs/ext/extension-api/src/contributors.rs
+++ b/codex-rs/ext/extension-api/src/contributors.rs
@@ -36,16 +36,17 @@ pub trait ContextContributor: Send + Sync {
 /// Implementations should use these callbacks to seed, rehydrate, or flush
 /// extension-private thread state. Heavy dependencies belong on the extension
 /// value created by the host, not in these inputs.
-pub trait ThreadLifecycleContributor<C>: Send + Sync {
+#[async_trait::async_trait]
+pub trait ThreadLifecycleContributor<C: Sync>: Send + Sync {
     /// Called after thread-scoped extension stores are created, before later
     /// contributors can read from them.
-    fn on_thread_start(&self, _input: ThreadStartInput<'_, C>) {}
+    async fn on_thread_start(&self, _input: ThreadStartInput<'_, C>) {}
 
     /// Called after the host constructs a runtime from persisted history.
-    fn on_thread_resume(&self, _input: ThreadResumeInput<'_>) {}
+    async fn on_thread_resume(&self, _input: ThreadResumeInput<'_>) {}
 
     /// Called before the host drops the thread runtime and thread-scoped store.
-    fn on_thread_stop(&self, _input: ThreadStopInput<'_>) {}
+    async fn on_thread_stop(&self, _input: ThreadStopInput<'_>) {}
 }
 
 /// Contributor for host-owned turn lifecycle gates.
@@ -53,16 +54,17 @@ pub trait ThreadLifecycleContributor<C>: Send + Sync {
 /// Implementations should use these callbacks to seed, observe, or clear
 /// extension-private turn state. The host exposes stable identifiers and
 /// extension stores instead of core runtime objects.
+#[async_trait::async_trait]
 pub trait TurnLifecycleContributor: Send + Sync {
     /// Called after turn-scoped extension stores are created, before the task
     /// for the turn starts running.
-    fn on_turn_start(&self, _input: TurnStartInput<'_>) {}
+    async fn on_turn_start(&self, _input: TurnStartInput<'_>) {}
 
     /// Called before the host drops the completed turn runtime and turn store.
-    fn on_turn_stop(&self, _input: TurnStopInput<'_>) {}
+    async fn on_turn_stop(&self, _input: TurnStopInput<'_>) {}
 
     /// Called after the host aborts a running turn.
-    fn on_turn_abort(&self, _input: TurnAbortInput<'_>) {}
+    async fn on_turn_abort(&self, _input: TurnAbortInput<'_>) {}
 }
 
 /// Contributor for host-owned configuration changes.
diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs
index 8193f5c1d1..f00e996315 100644
--- a/codex-rs/ext/extension-api/src/registry.rs
+++ b/codex-rs/ext/extension-api/src/registry.rs
@@ -12,7 +12,7 @@ use crate::TurnItemContributor;
 use crate::TurnLifecycleContributor;
 
 /// Mutable registry used while hosts register typed runtime contributions.
-pub struct ExtensionRegistryBuilder<C> {
+pub struct ExtensionRegistryBuilder<C: Sync> {
     thread_lifecycle_contributors: Vec<Arc<dyn ThreadLifecycleContributor<C>>>,
     turn_lifecycle_contributors: Vec<Arc<dyn TurnLifecycleContributor>>,
     config_contributors: Vec<Arc<dyn ConfigContributor<C>>>,
@@ -23,7 +23,7 @@ pub struct ExtensionRegistryBuilder<C> {
     approval_review_contributors: Vec<Arc<dyn ApprovalReviewContributor>>,
 }
 
-impl<C> Default for ExtensionRegistryBuilder<C> {
+impl<C: Sync> Default for ExtensionRegistryBuilder<C> {
     fn default() -> Self {
         Self {
             thread_lifecycle_contributors: Vec::new(),
@@ -38,7 +38,7 @@ impl<C> Default for ExtensionRegistryBuilder<C> {
     }
 }
 
-impl<C> ExtensionRegistryBuilder<C> {
+impl<C: Sync> ExtensionRegistryBuilder<C> {
     /// Creates an empty registry builder.
     pub fn new() -> Self {
         Self::default()
@@ -103,7 +103,7 @@ impl<C> ExtensionRegistryBuilder<C> {
 }
 
 /// Immutable typed registry produced after extensions are installed.
-pub struct ExtensionRegistry<C> {
+pub struct ExtensionRegistry<C: Sync> {
     thread_lifecycle_contributors: Vec<Arc<dyn ThreadLifecycleContributor<C>>>,
     turn_lifecycle_contributors: Vec<Arc<dyn TurnLifecycleContributor>>,
     config_contributors: Vec<Arc<dyn ConfigContributor<C>>>,
@@ -114,7 +114,7 @@ pub struct ExtensionRegistry<C> {
     approval_review_contributors: Vec<Arc<dyn ApprovalReviewContributor>>,
 }
 
-impl<C> ExtensionRegistry<C> {
+impl<C: Sync> ExtensionRegistry<C> {
     /// Returns the registered thread-lifecycle contributors.
     pub fn thread_lifecycle_contributors(&self) -> &[Arc<dyn ThreadLifecycleContributor<C>>] {
         &self.thread_lifecycle_contributors
@@ -165,6 +165,6 @@ impl<C> ExtensionRegistry<C> {
 }
 
 /// Creates an empty shared registry for hosts that do not register contributions.
-pub fn empty_extension_registry<C>() -> Arc<ExtensionRegistry<C>> {
+pub fn empty_extension_registry<C: Sync>() -> Arc<ExtensionRegistry<C>> {
     Arc::new(ExtensionRegistryBuilder::new().build())
 }
diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs
index 0d78742962..97d7157298 100644
--- a/codex-rs/ext/goal/src/extension.rs
+++ b/codex-rs/ext/goal/src/extension.rs
@@ -102,11 +102,12 @@ fn missing_backend_message() -> String {
     "goal tools are not connected to host goal persistence yet".to_string()
 }
 
+#[async_trait]
 impl<C> ThreadLifecycleContributor<C> for GoalExtension<C>
 where
     C: Send + Sync + 'static,
 {
-    fn on_thread_start(&self, input: ThreadStartInput<'_, C>) {
+    async fn on_thread_start(&self, input: ThreadStartInput<'_, C>) {
         input
             .thread_store
             .insert(GoalExtensionConfig::from_enabled((self.goals_enabled)(
@@ -135,11 +136,12 @@ where
     }
 }
 
+#[async_trait]
 impl<C> TurnLifecycleContributor for GoalExtension<C>
 where
     C: Send + Sync + 'static,
 {
-    fn on_turn_start(&self, input: TurnStartInput<'_>) {
+    async fn on_turn_start(&self, input: TurnStartInput<'_>) {
         if !goal_enabled(input.thread_store) {
             return;
         }
@@ -150,7 +152,7 @@ where
         accounting_state(input.thread_store).start_turn(input.turn_store.level_id());
     }
 
-    fn on_turn_stop(&self, input: TurnStopInput<'_>) {
+    async fn on_turn_stop(&self, input: TurnStopInput<'_>) {
         if !goal_enabled(input.thread_store) {
             return;
         }
@@ -164,7 +166,7 @@ where
         accounting_state(input.thread_store).stop_turn(input.turn_store.level_id());
     }
 
-    fn on_turn_abort(&self, input: TurnAbortInput<'_>) {
+    async fn on_turn_abort(&self, input: TurnAbortInput<'_>) {
         if !goal_enabled(input.thread_store) {
             return;
         }
diff --git a/codex-rs/ext/guardian/Cargo.toml b/codex-rs/ext/guardian/Cargo.toml
index 513254b7cf..53e553ec1c 100644
--- a/codex-rs/ext/guardian/Cargo.toml
+++ b/codex-rs/ext/guardian/Cargo.toml
@@ -14,6 +14,7 @@ doctest = false
 workspace = true
 
 [dependencies]
+async-trait = { workspace = true }
 codex-core = { workspace = true }
 codex-extension-api = { workspace = true }
 codex-protocol = { workspace = true }
diff --git a/codex-rs/ext/guardian/src/lib.rs b/codex-rs/ext/guardian/src/lib.rs
index 809a381471..0591887c25 100644
--- a/codex-rs/ext/guardian/src/lib.rs
+++ b/codex-rs/ext/guardian/src/lib.rs
@@ -47,11 +47,12 @@ impl GuardianThreadContext {
     }
 }
 
+#[async_trait::async_trait]
 impl<S> ThreadLifecycleContributor<Config> for GuardianExtension<S>
 where
     S: Send + Sync,
 {
-    fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) {
+    async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) {
         let Ok(forked_from_thread_id) = ThreadId::from_string(input.thread_store.level_id()) else {
             return;
         };
diff --git a/codex-rs/ext/memories/src/extension.rs b/codex-rs/ext/memories/src/extension.rs
index 1a52eb2cb9..a50949e919 100644
--- a/codex-rs/ext/memories/src/extension.rs
+++ b/codex-rs/ext/memories/src/extension.rs
@@ -58,8 +58,9 @@ impl ContextContributor for MemoriesExtension {
     }
 }
 
+#[async_trait::async_trait]
 impl ThreadLifecycleContributor<Config> for MemoriesExtension {
-    fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) {
+    async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) {
         input
             .thread_store
             .insert(MemoriesExtensionConfig::from_config(input.config));