ok

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<Self>, 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<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,
+        _context: &(),
+        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)
+}