diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs b/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs new file mode 100644 index 0000000000..e13dfbe8db --- /dev/null +++ b/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs @@ -0,0 +1,102 @@ +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_extension_api::CodexExtension; +use codex_extension_api::ContextContributor; +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::PromptFragment; + +/// Small tutorial extension that installs two prompt contributors. +#[derive(Debug, Default)] +pub struct SharedStateExtension; + +impl CodexExtension<()> for SharedStateExtension { + fn install(self: Arc, registry: &mut ExtensionRegistryBuilder<()>) { + registry.prompt_contributor(Arc::new(StyleContributor)); + registry.prompt_contributor(Arc::new(UsageContributor)); + } +} + +#[derive(Debug)] +struct StyleContributor; + +impl ContextContributor<()> for StyleContributor { + fn contribute( + &self, + _context: &(), + session_store: &ExtensionData, + thread_store: &ExtensionData, + ) -> Vec { + contribution_counts(session_store).record_style(); + contribution_counts(thread_store).record_style(); + + vec![PromptFragment::developer_policy( + "Prefer short answers unless the user asks for detail.", + )] + } +} + +#[derive(Debug)] +struct UsageContributor; + +impl ContextContributor<()> for UsageContributor { + fn contribute( + &self, + _context: &(), + session_store: &ExtensionData, + thread_store: &ExtensionData, + ) -> Vec { + contribution_counts(session_store).record_usage(); + contribution_counts(thread_store).record_usage(); + + vec![PromptFragment::developer_capability( + "This extension can contribute more than one prompt fragment.", + )] + } +} + +/// Returns how many style contributions were recorded in `store`. +pub fn recorded_style_contributions(store: &ExtensionData) -> u64 { + store + .get::() + .map(|counts| counts.style()) + .unwrap_or_default() +} + +/// Returns how many usage contributions were recorded in `store`. +pub fn recorded_usage_contributions(store: &ExtensionData) -> u64 { + store + .get::() + .map(|counts| counts.usage()) + .unwrap_or_default() +} + +#[derive(Debug, Default)] +struct ContributionCounts { + style: AtomicU64, + usage: AtomicU64, +} + +impl ContributionCounts { + fn record_style(&self) { + self.style.fetch_add(1, Ordering::Relaxed); + } + + fn record_usage(&self) { + self.usage.fetch_add(1, Ordering::Relaxed); + } + + fn style(&self) -> u64 { + self.style.load(Ordering::Relaxed) + } + + fn usage(&self) -> u64 { + self.usage.load(Ordering::Relaxed) + } +} + +fn contribution_counts(store: &ExtensionData) -> Arc { + store.get_or_init::(Default::default) +}