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))); }