Codex Handbook
ext/extension-api/examples/enabled_extensions/shared_state_extension.rs 98 lines
use std::sync::Arc;use std::sync::atomic::AtomicU64;use std::sync::atomic::Ordering;use codex_extension_api::ContextContributor;use codex_extension_api::ExtensionData;use codex_extension_api::ExtensionRegistryBuilder;use codex_extension_api::PromptFragment;/// Installs the tutorial contributors used by the example host.pub fn install(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<'a>(        &'a self,        session_store: &'a ExtensionData,        thread_store: &'a ExtensionData,    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {        Box::pin(async move {            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<'a>(        &'a self,        session_store: &'a ExtensionData,        thread_store: &'a ExtensionData,    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {        Box::pin(async move {            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::<ContributionCounts>()        .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::<ContributionCounts>()        .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<ContributionCounts> {    store.get_or_init::<ContributionCounts>(Default::default)}