From 5ab7e6b4c68bf6d19025c95692da89b170125c68 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 13 May 2026 13:11:30 +0200 Subject: [PATCH] feat: add thread lifecycle contributor hooks (#22476) ## Why Extensions that need thread-scoped state currently only get a start-time callback. That is enough for seeding stores, but it leaves the host without a shared extension seam for later thread rehydrate and flush work as thread ownership evolves. This PR turns that start-only seam into a host-owned thread lifecycle contributor contract so extension-private state can stay behind the extension API instead of leaking extra orchestration through core. ## What changed - Replaced `ThreadStartContributor` with `ThreadLifecycleContributor` and added typed lifecycle inputs for thread start, resume, and stop. The contract lives in [`contributors/thread_lifecycle.rs`](https://github.com/openai/codex/blob/d0e9211f70e58d6b07ef07e84f359d1b9aa25955/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs#L1-L64). - Kept the existing start-time behavior intact by routing session construction through `on_thread_start`. - Invoked `on_thread_stop` during session shutdown before thread-scoped extension state is dropped, while isolating contributor failures behind warning logs. - Migrated `git-attribution` and `guardian` onto the lifecycle registration path. - Renamed the extension registry plumbing from start-specific contributors to lifecycle-specific contributors. ## Notes `on_thread_resume` is introduced at the API boundary here so extensions can target the final lifecycle shape; host resume dispatch can be wired where that runtime path is finalized. --- codex-rs/Cargo.lock | 1 - codex-rs/core/src/codex_thread.rs | 16 +++++++++ codex-rs/core/src/session/handlers.rs | 8 +++++ codex-rs/core/src/session/session.rs | 12 +++---- codex-rs/core/src/thread_manager.rs | 9 ++--- .../ext/extension-api/src/contributors.rs | 34 +++++++++++------- .../src/contributors/thread_lifecycle.rs | 35 +++++++++++++++++++ codex-rs/ext/extension-api/src/lib.rs | 5 ++- codex-rs/ext/extension-api/src/registry.rs | 25 +++++++------ codex-rs/ext/git-attribution/Cargo.toml | 1 - codex-rs/ext/git-attribution/src/lib.rs | 22 +++++------- codex-rs/ext/guardian/src/lib.rs | 20 ++++------- 12 files changed, 124 insertions(+), 64 deletions(-) create mode 100644 codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c2b95f5a9..cd045e1f9 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2918,7 +2918,6 @@ dependencies = [ "codex-core", "codex-extension-api", "codex-features", - "codex-protocol", "pretty_assertions", ] diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index c110d30f6..60b77067a 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -140,6 +140,22 @@ impl CodexThread { self.codex.session_loop_termination.clone().await; } + pub(crate) fn emit_thread_resume_lifecycle(&self) { + for contributor in self + .codex + .session + .services + .extensions + .thread_lifecycle_contributors() + { + contributor.on_thread_resume(codex_extension_api::ThreadResumeInput { + thread_id: self.codex.session.conversation_id, + session_store: &self.codex.session.services.session_extension_data, + thread_store: &self.codex.session.services.thread_extension_data, + }); + } + } + pub async fn apply_goal_resume_runtime_effects(&self) -> anyhow::Result<()> { self.codex .session diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 03f29bac5..a3de362d6 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -643,6 +643,14 @@ pub async fn shutdown(sess: &Arc, 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, + }); + } + // Gracefully flush and shutdown thread persistence on session end so tests // that inspect durable state do not race with the background writer. if let Some(live_thread) = sess.live_thread() diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 747a6f5b4..6ca8ac1bf 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -813,13 +813,13 @@ impl Session { let agent_control = agent_control.with_session_id(session_id); let session_extension_data = codex_extension_api::ExtensionData::new(); let thread_extension_data = codex_extension_api::ExtensionData::new(); - for contributor in extensions.thread_start_contributors() { - contributor.contribute( + for contributor in extensions.thread_lifecycle_contributors() { + contributor.on_thread_start(codex_extension_api::ThreadStartInput { thread_id, - config.as_ref(), - &session_extension_data, - &thread_extension_data, - ); + config: config.as_ref(), + session_store: &session_extension_data, + thread_store: &thread_extension_data, + }); } let services = SessionServices { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index ce588d0a9..7b3d0a992 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1227,10 +1227,11 @@ impl ThreadManagerState { let new_thread = self .finalize_thread_spawn(codex, thread_id, tracked_session_source) .await?; - if is_resumed_thread - && let Err(err) = new_thread.thread.apply_goal_resume_runtime_effects().await - { - warn!("failed to apply goal resume runtime effects: {err}"); + if is_resumed_thread { + new_thread.thread.emit_thread_resume_lifecycle(); + if let Err(err) = new_thread.thread.apply_goal_resume_runtime_effects().await { + warn!("failed to apply goal resume runtime effects: {err}"); + } } Ok(new_thread) } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 99505e2af..ce8442d0e 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -1,33 +1,24 @@ use std::future::Future; use std::sync::Arc; -use codex_protocol::ThreadId; use codex_protocol::items::TurnItem; use codex_protocol::protocol::ReviewDecision; use crate::ExtensionData; mod prompt; +mod thread_lifecycle; mod tools; pub use prompt::PromptFragment; pub use prompt::PromptSlot; +pub use thread_lifecycle::ThreadResumeInput; +pub use thread_lifecycle::ThreadStartInput; +pub use thread_lifecycle::ThreadStopInput; pub use tools::ExtensionToolExecutor; pub use tools::ExtensionToolFuture; pub use tools::ExtensionToolOutput; -/// Contributor that receives the live thread id and host-owned thread-start -/// input before later contributors read from extension stores. -pub trait ThreadStartContributor: Send + Sync { - fn contribute( - &self, - thread_id: ThreadId, - input: &C, - session_store: &ExtensionData, - thread_store: &ExtensionData, - ); -} - /// Extension contribution that adds prompt fragments during prompt assembly. pub trait ContextContributor: Send + Sync { fn contribute( @@ -37,6 +28,23 @@ pub trait ContextContributor: Send + Sync { ) -> Vec; } +/// Contributor for host-owned thread lifecycle gates. +/// +/// 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: 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>) {} + + /// Called after the host constructs a runtime from persisted history. + 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<'_>) {} +} + /// Extension contribution that exposes native tools owned by a feature. pub trait ToolContributor: Send + Sync { /// Returns the native tools visible for the supplied extension stores. diff --git a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs new file mode 100644 index 000000000..41da56c03 --- /dev/null +++ b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs @@ -0,0 +1,35 @@ +use codex_protocol::ThreadId; + +use crate::ExtensionData; + +/// Input supplied when the host starts a runtime for a thread. +pub struct ThreadStartInput<'a, C> { + /// Identifier for the thread whose runtime is starting. + pub thread_id: ThreadId, + /// Host configuration visible at thread start. + pub config: &'a C, + /// Store scoped to the host session runtime. + pub session_store: &'a ExtensionData, + /// Store scoped to this thread runtime. + pub thread_store: &'a ExtensionData, +} + +/// Input supplied when the host resumes an existing thread. +pub struct ThreadResumeInput<'a> { + /// Identifier for the thread being resumed. + pub thread_id: ThreadId, + /// Store scoped to the host session runtime. + pub session_store: &'a ExtensionData, + /// Store scoped to this thread runtime. + pub thread_store: &'a ExtensionData, +} + +/// Input supplied when the host stops a thread runtime. +pub struct ThreadStopInput<'a> { + /// Identifier for the thread whose runtime is stopping. + pub thread_id: ThreadId, + /// Store scoped to the host session runtime. + pub session_store: &'a ExtensionData, + /// Store scoped to this thread runtime. + pub thread_store: &'a ExtensionData, +} diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index f1667b777..dab1ee46f 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -21,7 +21,10 @@ pub use contributors::ExtensionToolFuture; pub use contributors::ExtensionToolOutput; pub use contributors::PromptFragment; pub use contributors::PromptSlot; -pub use contributors::ThreadStartContributor; +pub use contributors::ThreadLifecycleContributor; +pub use contributors::ThreadResumeInput; +pub use contributors::ThreadStartInput; +pub use contributors::ThreadStopInput; pub use contributors::ToolContributor; pub use contributors::TurnItemContributionFuture; pub use contributors::TurnItemContributor; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index 4108ee492..1793c9ed0 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -4,13 +4,13 @@ use crate::ApprovalReviewContributor; use crate::ApprovalReviewFuture; use crate::ContextContributor; use crate::ExtensionData; -use crate::ThreadStartContributor; +use crate::ThreadLifecycleContributor; use crate::ToolContributor; use crate::TurnItemContributor; /// Mutable registry used while hosts register typed runtime contributions. pub struct ExtensionRegistryBuilder { - thread_start_contributors: Vec>>, + thread_lifecycle_contributors: Vec>>, context_contributors: Vec>, tool_contributors: Vec>, turn_item_contributors: Vec>, @@ -20,7 +20,7 @@ pub struct ExtensionRegistryBuilder { impl Default for ExtensionRegistryBuilder { fn default() -> Self { Self { - thread_start_contributors: Vec::new(), + thread_lifecycle_contributors: Vec::new(), approval_review_contributors: Vec::new(), context_contributors: Vec::new(), tool_contributors: Vec::new(), @@ -40,9 +40,12 @@ impl ExtensionRegistryBuilder { self.approval_review_contributors.push(contributor); } - /// Registers one thread-start contributor. - pub fn thread_start_contributor(&mut self, contributor: Arc>) { - self.thread_start_contributors.push(contributor); + /// Registers one thread-lifecycle contributor. + pub fn thread_lifecycle_contributor( + &mut self, + contributor: Arc>, + ) { + self.thread_lifecycle_contributors.push(contributor); } /// Registers one prompt contributor. @@ -63,7 +66,7 @@ impl ExtensionRegistryBuilder { /// Finishes construction and returns the immutable registry. pub fn build(self) -> ExtensionRegistry { ExtensionRegistry { - thread_start_contributors: self.thread_start_contributors, + thread_lifecycle_contributors: self.thread_lifecycle_contributors, approval_review_contributors: self.approval_review_contributors, context_contributors: self.context_contributors, tool_contributors: self.tool_contributors, @@ -74,7 +77,7 @@ impl ExtensionRegistryBuilder { /// Immutable typed registry produced after extensions are installed. pub struct ExtensionRegistry { - thread_start_contributors: Vec>>, + thread_lifecycle_contributors: Vec>>, context_contributors: Vec>, tool_contributors: Vec>, turn_item_contributors: Vec>, @@ -82,9 +85,9 @@ pub struct ExtensionRegistry { } impl ExtensionRegistry { - /// Returns the registered thread-start contributors. - pub fn thread_start_contributors(&self) -> &[Arc>] { - &self.thread_start_contributors + /// Returns the registered thread-lifecycle contributors. + pub fn thread_lifecycle_contributors(&self) -> &[Arc>] { + &self.thread_lifecycle_contributors } /// Claims the first rendered approval-review prompt accepted by an diff --git a/codex-rs/ext/git-attribution/Cargo.toml b/codex-rs/ext/git-attribution/Cargo.toml index 500ce47ae..edf8c0719 100644 --- a/codex-rs/ext/git-attribution/Cargo.toml +++ b/codex-rs/ext/git-attribution/Cargo.toml @@ -16,7 +16,6 @@ workspace = true codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } -codex-protocol = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/ext/git-attribution/src/lib.rs b/codex-rs/ext/git-attribution/src/lib.rs index eae906fa5..1df7815d6 100644 --- a/codex-rs/ext/git-attribution/src/lib.rs +++ b/codex-rs/ext/git-attribution/src/lib.rs @@ -5,9 +5,9 @@ use codex_extension_api::ContextContributor; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::PromptFragment; -use codex_extension_api::ThreadStartContributor; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStartInput; use codex_features::Feature; -use codex_protocol::ThreadId; const DEFAULT_ATTRIBUTION_VALUE: &str = "Codex "; @@ -40,17 +40,11 @@ struct GitAttributionConfig { prompt: Option, } -impl ThreadStartContributor for GitAttributionExtension { - fn contribute( - &self, - _thread_id: ThreadId, - config: &Config, - _session_store: &ExtensionData, - thread_store: &ExtensionData, - ) { - thread_store.insert(GitAttributionConfig { - enabled: config.features.enabled(Feature::CodexGitCommit), - prompt: config.commit_attribution.clone(), +impl ThreadLifecycleContributor for GitAttributionExtension { + fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { + input.thread_store.insert(GitAttributionConfig { + enabled: input.config.features.enabled(Feature::CodexGitCommit), + prompt: input.config.commit_attribution.clone(), }); } } @@ -58,7 +52,7 @@ impl ThreadStartContributor for GitAttributionExtension { /// Installs the git-attribution contributors into the extension registry. pub fn install(registry: &mut ExtensionRegistryBuilder) { let extension = Arc::new(GitAttributionExtension); - registry.thread_start_contributor(extension.clone()); + registry.thread_lifecycle_contributor(extension.clone()); registry.prompt_contributor(extension); } diff --git a/codex-rs/ext/guardian/src/lib.rs b/codex-rs/ext/guardian/src/lib.rs index d8c9834d6..1754f6e64 100644 --- a/codex-rs/ext/guardian/src/lib.rs +++ b/codex-rs/ext/guardian/src/lib.rs @@ -3,9 +3,9 @@ use std::sync::Arc; use codex_core::config::Config; use codex_extension_api::AgentSpawnFuture; use codex_extension_api::AgentSpawner; -use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistryBuilder; -use codex_extension_api::ThreadStartContributor; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStartInput; use codex_protocol::ThreadId; /// Guardian extension dependencies supplied by the host at construction time. @@ -47,19 +47,13 @@ impl GuardianThreadContext { } } -impl ThreadStartContributor for GuardianExtension +impl ThreadLifecycleContributor for GuardianExtension where S: Send + Sync, { - fn contribute( - &self, - thread_id: ThreadId, - _input: &Config, - _session_store: &ExtensionData, - thread_store: &ExtensionData, - ) { - thread_store.insert(GuardianThreadContext { - forked_from_thread_id: thread_id, + fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { + input.thread_store.insert(GuardianThreadContext { + forked_from_thread_id: input.thread_id, }); } } @@ -69,5 +63,5 @@ pub fn install(registry: &mut ExtensionRegistryBuilder, agent_spawner where S: Send + Sync + 'static, { - registry.thread_start_contributor(Arc::new(GuardianExtension::new(agent_spawner))); + registry.thread_lifecycle_contributor(Arc::new(GuardianExtension::new(agent_spawner))); }