codex/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs

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(
        &self,
        session_store: &ExtensionData,
        thread_store: &ExtensionData,
    ) -> Vec<PromptFragment> {
        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,
        session_store: &ExtensionData,
        thread_store: &ExtensionData,
    ) -> Vec<PromptFragment> {
        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)
}