Add state managemetn

codex-rs/ext/extension-api/Cargo.toml

@@ -17,7 +17,5 @@ serde_json = { workspace = true }
 thiserror = { workspace = true }
 
 [dev-dependencies]
-codex-guardian = { workspace = true }
 codex-git-attribution = { workspace = true }
 codex-memories = { workspace = true }
-codex-multi-agent-v2 = { workspace = true }

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

@@ -1,112 +1,89 @@
 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;
+use codex_extension_api::Stores;
+use codex_extension_api::scopes::Thread;
+use codex_extension_api::scopes::Turn;
 use codex_git_attribution as git_attribution;
-use codex_guardian as guardian;
 use codex_memories::MemoriesExtension;
-use codex_multi_agent_v2 as multi_agent_v2;
-use codex_multi_agent_v2::UsageHintAudience;
 
 fn main() {
+    // 1. Build the concrete extension values owned by the host.
+    let memories = Arc::new(MemoriesExtension::new(
+        Some("Use memories when they help answer the user.".to_string()),
+        std::env::temp_dir().join("codex-memories-example"),
+    ));
+    let git_attribution = git_attribution::extension();
+
+    // 2. Install those extensions into one registry for the runtime context type
+    //    this host exposes to contributors.
     let registry = ExtensionRegistryBuilder::<ctx::RuntimeContext>::new()
-        .with_extension(guardian::extension())
-        .with_extension(git_attribution::extension())
-        .with_extension(Arc::new(MemoriesExtension::with_read_prompt(
-            "Please use FS access bla bla bla.".to_string(),
-            std::env::temp_dir().join("codex-memories-example"),
-        )))
-        .with_extension(multi_agent_v2::extension())
+        .with_extension(git_attribution)
+        .with_extension(memories)
         .build();
 
-    let root_context = ctx::RuntimeContext {
-        automatic_review_enabled: true,
-        approval_policy_allows_automatic_review: true,
-        is_guardian_reviewer: false,
-        guardian_policy_prompt: Some("Guardian policy.".to_string()),
+    // 3. Build the runtime facts contributors can inspect when they decide
+    //    whether they are active for this request.
+    //
+    // Ideally, this is instead the TurnContext or the Config or so
+    let context = ctx::RuntimeContext {
         commit_attribution: None,
         memory_tool_enabled: true,
         use_memories: true,
-        multi_agent_v2_enabled: true,
-        multi_agent_v2_usage_hint_audience: UsageHintAudience::RootAgent,
-        root_agent_usage_hint_text: Some("Root-agent usage hint.".to_string()),
-        subagent_usage_hint_text: Some("Subagent usage hint.".to_string()),
-    };
-    let memories_disabled_context = ctx::RuntimeContext {
-        use_memories: false,
-        ..root_context.clone()
     };
 
-    // Build the prompt
+    // 4. Assemble the stores that exist at this insertion point.
+    //
+    //    Another insertion point could expose only `Thread`, or could add more
+    //    scopes later. Contributors receive the same dynamic bag and address the
+    //    scopes they need by marker type. The idea would be to have some re-usable contributors
+    //    such as the Output one, but I'm happy to negociate this one
+    let thread_data = ExtensionData::new();
+    let turn_data = ExtensionData::new();
+    let stores = codex_extension_api::stores! {
+        Thread => &thread_data,
+        Turn => &turn_data,
+    };
+
+    // 5. Invoke whichever contribution families this insertion point needs.
     let prompt_fragments = registry
         .context_contributors()
         .iter()
-        .flat_map(|contributor| contributor.contribute(&root_context))
+        .flat_map(|contributor| contributor.contribute(&context, &stores))
         .collect::<Vec<_>>();
 
-    // Get native tools
     let tools = registry
         .tool_contributors()
         .iter()
-        .flat_map(|contributor| contributor.tools(&root_context))
+        .flat_map(|contributor| contributor.tools(&context, &stores))
         .collect::<Vec<_>>();
-    let tools_without_memories = registry
-        .tool_contributors()
-        .iter()
-        .flat_map(|contributor| contributor.tools(&memories_disabled_context))
-        .collect::<Vec<_>>();
-    let active_approval_interceptors = registry
-        .approval_interceptor_contributors()
-        .iter()
-        .filter(|contributor| contributor.intercepts_approvals(&root_context))
-        .count();
 
     println!("prompt fragments: {}", prompt_fragments.len());
-    println!("approval interceptors: {active_approval_interceptors}");
     println!("native tools: {}", tools.len());
-    println!(
-        "native tools when use_memories=false: {}",
-        tools_without_memories.len()
-    );
 }
+
+// Just for the machinerie such that it compiles
+#[derive(Debug, Default)]
+struct ThreadPromptStats {
+    prompt_builds: AtomicU64,
+}
+
 mod ctx {
     use codex_git_attribution::GitAttributionContext;
-    use codex_guardian::GuardianContext;
     use codex_memories::ctx::MemoriesContext;
-    use codex_multi_agent_v2::MultiAgentV2Context;
-    use codex_multi_agent_v2::UsageHintAudience;
 
+    /// Host-owned projection exposed to extension contributors.
     #[derive(Clone)]
     pub struct RuntimeContext {
-        pub automatic_review_enabled: bool,
-        pub approval_policy_allows_automatic_review: bool,
-        pub is_guardian_reviewer: bool,
-        pub guardian_policy_prompt: Option<String>,
-        // Ideally this should be at the config layer instead
         pub commit_attribution: Option<String>,
         pub memory_tool_enabled: bool,
         pub use_memories: bool,
-        pub multi_agent_v2_enabled: bool,
-        pub multi_agent_v2_usage_hint_audience: UsageHintAudience,
-        pub root_agent_usage_hint_text: Option<String>,
-        pub subagent_usage_hint_text: Option<String>,
-    }
-
-    impl GuardianContext for RuntimeContext {
-        fn automatic_review_enabled(&self) -> bool {
-            self.automatic_review_enabled
-        }
-
-        fn approval_policy_allows_automatic_review(&self) -> bool {
-            self.approval_policy_allows_automatic_review
-        }
-
-        fn is_guardian_reviewer(&self) -> bool {
-            self.is_guardian_reviewer
-        }
-
-        fn guardian_policy_prompt(&self) -> Option<&str> {
-            self.guardian_policy_prompt.as_deref()
-        }
     }
 
     impl GitAttributionContext for RuntimeContext {
@@ -124,22 +101,4 @@ mod ctx {
             self.use_memories
         }
     }
-
-    impl MultiAgentV2Context for RuntimeContext {
-        fn multi_agent_v2_enabled(&self) -> bool {
-            self.multi_agent_v2_enabled
-        }
-
-        fn usage_hint_audience(&self) -> UsageHintAudience {
-            self.multi_agent_v2_usage_hint_audience
-        }
-
-        fn root_agent_usage_hint_text(&self) -> Option<&str> {
-            self.root_agent_usage_hint_text.as_deref()
-        }
-
-        fn subagent_usage_hint_text(&self) -> Option<&str> {
-            self.subagent_usage_hint_text.as_deref()
-        }
-    }
 }

codex-rs/ext/extension-api/src/contributors.rs

@@ -1,14 +1,8 @@
-//! Runtime contribution ports exposed by the extension API.
-//!
-//! Extensions can install several independent contribution types. This module
-//! is the small public façade over those ports: callers can import the common
-//! vocabulary from one place, while each contribution family keeps its own
-//! supporting types nearby.
-
 use std::future::Future;
 
+use crate::Stores;
+
 mod prompt;
-mod session_lifecycle;
 mod tool;
 
 pub use prompt::PromptFragment;
@@ -18,15 +12,14 @@ pub use tool::ToolContribution;
 pub use tool::ToolHandler;
 
 /// Extension contribution that adds prompt fragments during prompt assembly.
-/// Arguably, can become async
 pub trait ContextContributor<C>: Send + Sync {
-    fn contribute(&self, context: &C) -> Vec<PromptFragment>; // TODO use existing fragments ofc
+    fn contribute(&self, context: &C, stores: &Stores<'_>) -> Vec<PromptFragment>;
 }
 
 /// Extension contribution that exposes native tools owned by a feature.
 pub trait ToolContributor<C>: Send + Sync {
     /// Returns the native tools visible for the supplied runtime context.
-    fn tools(&self, context: &C) -> Vec<ToolContribution<C>>;
+    fn tools(&self, context: &C, stores: &Stores<'_>) -> Vec<ToolContribution<C>>;
 }
 
 /// Future returned by one ordered output contribution.
@@ -35,24 +28,21 @@ pub type OutputContributionFuture<'a> =
 
 /// Ordered post-processing contribution for one completed output value.
 ///
-/// Implementations may inspect or mutate `output`; hosts are expected to run
-/// contributors sequentially so each contributor observes the result of the
-/// previous one.
+/// This is kept abstract so that we can re-use it at multiple places. I kind of like it but I understand
+/// if this is problematic with bad devs
 pub trait OutputContributor<C, O>: Send + Sync {
-    fn contribute<'a>(&'a self, context: &'a C, output: &'a mut O) -> OutputContributionFuture<'a>;
+    fn contribute<'a>(
+        &'a self,
+        context: &'a C,
+        stores: &'a Stores<'a>,
+        output: &'a mut O,
+    ) -> OutputContributionFuture<'a>;
 }
 
-// TODO: WIP
+// TODO: WIP (do not consider)
 /// Extension contribution that can claim approval requests for a runtime context.
 /// (ideally we can replace it by a session lifecycle thing or a request contributor?)
 pub trait ApprovalInterceptorContributor<C>: Send + Sync {
     /// Returns whether this contributor should intercept approvals in `context`.
-    fn intercepts_approvals(&self, context: &C) -> bool;
-}
-
-pub trait SessionLifecycleContributor<S, T>: Send + Sync {
-    fn on_lifecycle_event(
-        &self,
-        event: session_lifecycle::Event<'_, S, T>,
-    ) -> impl Future<Output = ()> + Send;
+    fn intercepts_approvals(&self, context: &C, stores: &Stores<'_>) -> bool;
 }

codex-rs/ext/extension-api/src/contributors/session_lifecycle.rs

@@ -1,6 +0,0 @@
-pub enum Event<'a, S, T> {
-    SessionStarted(&'a S),
-    TurnStarted(&'a T),
-    TurnFinished { turn: &'a T, outcome: String }, // Make outcome better?!
-    SessionStopping(&'a S),
-}

codex-rs/ext/extension-api/src/lib.rs

@@ -18,5 +18,4 @@ pub use registry::ExtensionRegistry;
 pub use registry::ExtensionRegistryBuilder;
 pub use state::ExtensionData;
 pub use state::Stores;
-pub use state::Thread;
-pub use state::Turn;
+pub use state::scopes;

codex-rs/ext/extension-api/src/state.rs

@@ -0,0 +1,207 @@
+use std::any::Any;
+use std::any::TypeId;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::PoisonError;
+
+type ErasedData = Arc<dyn Any + Send + Sync>;
+
+/// Built-in extension-store scope markers.
+pub mod scopes {
+    pub enum Thread {}
+    pub enum Turn {}
+}
+
+/// Typed extension-owned data attached to one host object.
+#[derive(Default, Debug)]
+pub struct ExtensionData {
+    entries: Mutex<HashMap<TypeId, ErasedData>>,
+}
+
+impl ExtensionData {
+    /// Creates an empty attachment map.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Returns the attached value of type `T`, if one exists.
+    pub fn get<T>(&self) -> Option<Arc<T>>
+    where
+        T: Any + Send + Sync,
+    {
+        let value = self.entries().get(&TypeId::of::<T>())?.clone();
+        Some(downcast_data(value))
+    }
+
+    /// Returns the attached value of type `T`, inserting one from `init` when absent.
+    ///
+    /// The initializer runs while this map is locked, so it should stay cheap;
+    /// heavyweight lazy work belongs inside the attached value itself.
+    pub fn get_or_init<T>(&self, init: impl FnOnce() -> T) -> Arc<T>
+    where
+        T: Any + Send + Sync,
+    {
+        let mut entries = self.entries();
+        let value = entries
+            .entry(TypeId::of::<T>())
+            .or_insert_with(|| Arc::new(init()));
+        downcast_data(Arc::clone(value))
+    }
+
+    /// Stores `value` as the attachment of type `T`, returning any previous value.
+    pub fn insert<T>(&self, value: T) -> Option<Arc<T>>
+    where
+        T: Any + Send + Sync,
+    {
+        self.entries()
+            .insert(TypeId::of::<T>(), Arc::new(value))
+            .map(downcast_data)
+    }
+
+    /// Removes and returns the attached value of type `T`, if one exists.
+    pub fn remove<T>(&self) -> Option<Arc<T>>
+    where
+        T: Any + Send + Sync,
+    {
+        self.entries().remove(&TypeId::of::<T>()).map(downcast_data)
+    }
+
+    fn entries(&self) -> std::sync::MutexGuard<'_, HashMap<TypeId, ErasedData>> {
+        self.entries.lock().unwrap_or_else(PoisonError::into_inner)
+    }
+}
+
+/// Dynamic set of host-owned extension stores visible at one contribution site.
+///
+/// The host decides which lifetime scopes are available for each insertion
+/// point. Contributors can then address one of those visible scopes by marker
+/// type while storing extension-private values inside the selected owner.
+#[derive(Default, Debug)]
+pub struct Stores<'a> {
+    stores: HashMap<TypeId, &'a ExtensionData>,
+}
+
+impl<'a> Stores<'a> {
+    /// Creates an empty set of visible stores.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Adds or replaces the store identified by scope marker `Scope`.
+    pub fn insert_scope<Scope: 'static>(
+        &mut self,
+        store: &'a ExtensionData,
+    ) -> Option<&'a ExtensionData> {
+        self.stores.insert(TypeId::of::<Scope>(), store)
+    }
+
+    /// Returns the value of type `T` from `Scope`, if one exists.
+    ///
+    /// # Panics
+    ///
+    /// Panics when `Scope` is unavailable at this contribution site. That is a
+    /// host/contributor wiring bug: insertion points define which scopes exist,
+    /// and contributors should only request scopes promised by their insertion point.
+    pub fn get<Scope, T>(&self) -> Option<Arc<T>>
+    where
+        Scope: 'static,
+        T: Any + Send + Sync,
+    {
+        self.store::<Scope>().get::<T>()
+    }
+
+    /// Returns the value of type `T`, inserting one into `Scope` when absent.
+    ///
+    /// # Panics
+    ///
+    /// Panics when `Scope` is unavailable at this contribution site. That is a
+    /// host/contributor wiring bug: insertion points define which scopes exist,
+    /// and contributors should only request scopes promised by their insertion point.
+    pub fn get_or_init<Scope, T>(&self, init: impl FnOnce() -> T) -> Arc<T>
+    where
+        Scope: 'static,
+        T: Any + Send + Sync,
+    {
+        self.store::<Scope>().get_or_init(init)
+    }
+
+    /// Returns the value of type `T`, inserting one into `Scope` when available.
+    ///
+    /// Use this when a contributor intentionally supports multiple insertion
+    /// points and can proceed without state from a scope that is not visible.
+    pub fn try_get_or_init<Scope, T>(&self, init: impl FnOnce() -> T) -> Option<Arc<T>>
+    where
+        Scope: 'static,
+        T: Any + Send + Sync,
+    {
+        self.try_store::<Scope>()
+            .map(|store| store.get_or_init(init))
+    }
+
+    /// Stores `value` in `Scope`, returning any previous value of type `T`.
+    ///
+    /// # Panics
+    ///
+    /// Panics when `Scope` is unavailable at this contribution site. That is a
+    /// host/contributor wiring bug: insertion points define which scopes exist,
+    /// and contributors should only request scopes promised by their insertion point.
+    pub fn insert<Scope, T>(&self, value: T) -> Option<Arc<T>>
+    where
+        Scope: 'static,
+        T: Any + Send + Sync,
+    {
+        self.store::<Scope>().insert(value)
+    }
+
+    /// Removes and returns the value of type `T` from `Scope`, if one exists.
+    ///
+    /// # Panics
+    ///
+    /// Panics when `Scope` is unavailable at this contribution site. That is a
+    /// host/contributor wiring bug: insertion points define which scopes exist,
+    /// and contributors should only request scopes promised by their insertion point.
+    pub fn remove<Scope, T>(&self) -> Option<Arc<T>>
+    where
+        Scope: 'static,
+        T: Any + Send + Sync,
+    {
+        self.store::<Scope>().remove::<T>()
+    }
+
+    fn try_store<Scope: 'static>(&self) -> Option<&ExtensionData> {
+        self.stores.get(&TypeId::of::<Scope>()).copied()
+    }
+
+    fn store<Scope: 'static>(&self) -> &ExtensionData {
+        self.try_store::<Scope>().unwrap_or_else(|| {
+            // This panic means a mistake made by a developer!!!
+            panic!(
+                "extension store for scope `{}` is unavailable",
+                std::any::type_name::<Scope>()
+            )
+        })
+    }
+}
+
+/// Builds a dynamic [`Stores`] bag from the scopes available at one insertion point.
+///
+/// I know people don't like macros but this is super cool IMO
+#[macro_export]
+macro_rules! stores {
+    ($($scope:ty => $store:expr),* $(,)?) => {{
+        let mut stores = $crate::Stores::new();
+        $(stores.insert_scope::<$scope>($store);)*
+        stores
+    }};
+}
+
+fn downcast_data<T>(value: ErasedData) -> Arc<T>
+where
+    T: Any + Send + Sync,
+{
+    let Ok(value) = value.downcast::<T>() else {
+        unreachable!("typed extension data stored an incompatible value");
+    };
+    value
+}

codex-rs/ext/git-attribution/src/lib.rs

@@ -1,21 +1,13 @@
-//! Git-attribution prompt contribution packaged as a Codex extension.
-
-#![forbid(unsafe_code)]
-
 use std::sync::Arc;
 
 use codex_extension_api::CodexExtension;
 use codex_extension_api::ContextContributor;
 use codex_extension_api::ExtensionRegistryBuilder;
 use codex_extension_api::PromptFragment;
+use codex_extension_api::Stores;
 
 const DEFAULT_ATTRIBUTION_VALUE: &str = "Codex <noreply@openai.com>";
 
-/// Runtime facts needed to render the git-attribution prompt contribution.
-///
-/// Hosts should return the current effective attribution override for the
-/// thread being assembled. `None` selects Codex's default attribution, while a
-/// blank string disables the contribution.
 pub trait GitAttributionContext {
     fn commit_attribution(&self) -> Option<&str>;
 }
@@ -40,7 +32,7 @@ impl GitAttributionExtension {
 }
 
 impl<C: GitAttributionContext> ContextContributor<C> for GitAttributionExtension {
-    fn contribute(&self, context: &C) -> Vec<PromptFragment> {
+    fn contribute(&self, context: &C, _stores: &Stores<'_>) -> Vec<PromptFragment> {
         self.instruction(context)
             .map(PromptFragment::developer_capability)
             .into_iter()
@@ -59,6 +51,7 @@ pub fn extension() -> Arc<GitAttributionExtension> {
     Arc::new(GitAttributionExtension::new())
 }
 
+// This is just a copy/paste
 fn build_commit_message_trailer(config_attribution: Option<&str>) -> Option<String> {
     let value = resolve_attribution_value(config_attribution)?;
     Some(format!("Co-authored-by: {value}"))

codex-rs/ext/guardian/BUILD.bazel

@@ -1,6 +0,0 @@
-load("//:defs.bzl", "codex_rust_crate")
-
-codex_rust_crate(
-    name = "guardian",
-    crate_name = "codex_guardian",
-)

codex-rs/ext/guardian/Cargo.toml

@@ -1,15 +0,0 @@
-[package]
-edition.workspace = true
-license.workspace = true
-name = "codex-guardian"
-version.workspace = true
-
-[lib]
-name = "codex_guardian"
-path = "src/lib.rs"
-
-[lints]
-workspace = true
-
-[dependencies]
-codex-extension-api = { workspace = true }

codex-rs/ext/guardian/src/lib.rs

@@ -1,75 +0,0 @@
-//! Guardian routing contribution packaged as a Codex extension.
-
-#![forbid(unsafe_code)]
-
-use std::sync::Arc;
-
-use codex_extension_api::ApprovalInterceptorContributor;
-use codex_extension_api::CodexExtension;
-use codex_extension_api::ContextContributor;
-use codex_extension_api::ExtensionRegistryBuilder;
-use codex_extension_api::PromptFragment;
-
-/// Runtime facts needed to expose Guardian surfaces.
-///
-/// Hosts should provide the effective approval settings for the current turn,
-/// whether the current session is the Guardian reviewer itself, and the prompt
-/// text to show to that reviewer.
-pub trait GuardianContext {
-    fn automatic_review_enabled(&self) -> bool;
-    fn approval_policy_allows_automatic_review(&self) -> bool;
-    fn is_guardian_reviewer(&self) -> bool;
-    fn guardian_policy_prompt(&self) -> Option<&str>;
-}
-
-/// Extension that contributes Guardian approval routing and reviewer policy.
-#[derive(Clone, Copy, Debug, Default)]
-pub struct GuardianExtension;
-
-impl GuardianExtension {
-    /// Creates an extension instance.
-    pub fn new() -> Self {
-        Self
-    }
-
-    /// Returns whether Guardian should intercept approvals in this context.
-    pub fn should_intercept_approvals<C: GuardianContext>(&self, context: &C) -> bool {
-        context.automatic_review_enabled() && context.approval_policy_allows_automatic_review()
-    }
-
-    /// Returns the policy prompt shown only to Guardian reviewer sessions.
-    pub fn policy_prompt<'a, C: GuardianContext>(&self, context: &'a C) -> Option<&'a str> {
-        if context.is_guardian_reviewer() {
-            context.guardian_policy_prompt()
-        } else {
-            None
-        }
-    }
-}
-
-impl<C: GuardianContext> ApprovalInterceptorContributor<C> for GuardianExtension {
-    fn intercepts_approvals(&self, context: &C) -> bool {
-        self.should_intercept_approvals(context)
-    }
-}
-
-impl<C: GuardianContext> ContextContributor<C> for GuardianExtension {
-    fn contribute(&self, context: &C) -> Vec<PromptFragment> {
-        self.policy_prompt(context)
-            .map(PromptFragment::separate_developer)
-            .into_iter()
-            .collect()
-    }
-}
-
-impl<C: GuardianContext> CodexExtension<C> for GuardianExtension {
-    fn install(self: Arc<Self>, registry: &mut ExtensionRegistryBuilder<C>) {
-        registry.approval_interceptor_contributor(self.clone());
-        registry.prompt_contributor(self);
-    }
-}
-
-/// Creates a shared Guardian extension instance.
-pub fn extension() -> Arc<GuardianExtension> {
-    Arc::new(GuardianExtension::new())
-}

codex-rs/ext/memories/Cargo.toml

@@ -22,3 +22,4 @@ serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["fs"] }
+tracing = { workspace = true }

codex-rs/ext/memories/src/citation_output.rs

@@ -1,5 +1,12 @@
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
+
 use codex_extension_api::OutputContributionFuture;
 use codex_extension_api::OutputContributor;
+use codex_extension_api::Stores;
+use codex_extension_api::scopes::Thread;
+use codex_extension_api::scopes::Turn;
 use codex_memories_read::citations::parse_memory_citation;
 use codex_protocol::items::AgentMessageContent;
 use codex_protocol::items::TurnItem;
@@ -7,10 +14,21 @@ use codex_utils_stream_parser::strip_citations;
 
 use crate::MemoriesExtension;
 
+#[derive(Debug, Default)]
+struct ThreadMemoriesState {
+    turns_with_citations: AtomicU64,
+}
+
+#[derive(Debug, Default)]
+struct TurnMemoriesState {
+    had_citation: AtomicBool,
+}
+
 impl<C> OutputContributor<C, TurnItem> for MemoriesExtension {
     fn contribute<'a>(
         &'a self,
         _context: &'a C,
+        stores: &'a Stores<'a>,
         output: &'a mut TurnItem,
     ) -> OutputContributionFuture<'a> {
         if let TurnItem::AgentMessage(agent_message) = output {
@@ -24,8 +42,29 @@ impl<C> OutputContributor<C, TurnItem> for MemoriesExtension {
             let (visible_text, citations) = strip_citations(&combined);
             agent_message.content = vec![AgentMessageContent::Text { text: visible_text }];
             agent_message.memory_citation = parse_memory_citation(citations);
+
+            if agent_message.memory_citation.is_some()
+                && let Some(turns_with_citations) = record_citation_seen(stores)
+            {
+                tracing::info!(turns_with_citations, "memory citation seen in turn");
+            }
         }
 
         Box::pin(std::future::ready(Ok(())))
     }
 }
+
+fn record_citation_seen(stores: &Stores<'_>) -> Option<u64> {
+    let turn_state = stores.get_or_init::<Turn, TurnMemoriesState>(Default::default);
+    if turn_state.had_citation.swap(true, Ordering::Relaxed) {
+        return None;
+    }
+
+    let thread_stats = stores.get_or_init::<Thread, ThreadMemoriesState>(Default::default);
+    Some(
+        thread_stats
+            .turns_with_citations
+            .fetch_add(1, Ordering::Relaxed)
+            + 1,
+    )
+}

codex-rs/ext/memories/src/lib.rs

@@ -1,21 +1,15 @@
-//! Memories tools and prompt contribution packaged as a Codex extension.
-
-#![forbid(unsafe_code)]
-
 mod citation_output;
 pub mod ctx;
 mod list_tool;
+mod prompt_contributor;
+mod tool_contributor;
 
 use std::path::PathBuf;
 use std::sync::Arc;
 
 use crate::ctx::MemoriesContext;
 use codex_extension_api::CodexExtension;
-use codex_extension_api::ContextContributor;
 use codex_extension_api::ExtensionRegistryBuilder;
-use codex_extension_api::PromptFragment;
-use codex_extension_api::ToolContribution;
-use codex_extension_api::ToolContributor;
 use codex_memories_read::build_memory_tool_developer_instructions;
 use codex_memories_read::memory_root;
 use codex_protocol::items::TurnItem;
@@ -26,33 +20,13 @@ use list_tool::ListMemoriesTool;
 #[derive(Clone, Debug)]
 pub struct MemoriesExtension {
     read_prompt: Option<String>,
-    list_tool: Arc<ListMemoriesTool>,
+    pub(crate) list_tool: Arc<ListMemoriesTool>, // This is just to have useful examples, it will disappear
 }
 
-impl<C: MemoriesContext + Send + Sync + 'static> ToolContributor<C> for MemoriesExtension {
-    fn tools(&self, context: &C) -> Vec<ToolContribution<C>> {
-        if !self.is_read_surface_enabled(context) {
-            return Vec::new();
-        }
-
-        vec![self.list_tool.contribution()]
-    }
-}
-
-impl<C: MemoriesContext> ContextContributor<C> for MemoriesExtension {
-    fn contribute(&self, context: &C) -> Vec<PromptFragment> {
-        if !self.is_read_surface_enabled(context) {
-            return Vec::new();
-        }
-
-        self.read_prompt()
-            .map(PromptFragment::developer_policy)
-            .into_iter()
-            .collect()
-    }
-}
-
-impl<C: MemoriesContext + Send + Sync + 'static> CodexExtension<C> for MemoriesExtension {
+impl<C> CodexExtension<C> for MemoriesExtension
+where
+    C: MemoriesContext + Send + Sync + 'static,
+{
     fn install(self: Arc<Self>, registry: &mut ExtensionRegistryBuilder<C>) {
         registry.tool_contributor(self.clone());
         registry.output_contributor::<TurnItem>(self.clone());
@@ -61,37 +35,29 @@ impl<C: MemoriesContext + Send + Sync + 'static> CodexExtension<C> for MemoriesE
 }
 
 impl MemoriesExtension {
-    fn new(read_prompt: Option<String>, memories_root: impl Into<PathBuf>) -> Self {
+    /// Creates a memories extension from the prompt text and backing directory
+    /// it should expose.
+    pub fn new(read_prompt: Option<String>, memories_root: impl Into<PathBuf>) -> Self {
         Self {
             read_prompt,
             list_tool: Arc::new(ListMemoriesTool::new(memories_root)),
         }
     }
 
-    /// Creates an extension with one pre-rendered read prompt and native tools.
-    pub fn with_read_prompt(read_prompt: String, memories_root: impl Into<PathBuf>) -> Self {
-        Self::new(Some(read_prompt), memories_root)
-    }
-
-    /// Creates an extension using the live memories read prompt for this Codex home.
-    pub async fn from_codex_home(codex_home: &AbsolutePathBuf) -> Self {
-        Self::new(
-            build_memory_tool_developer_instructions(codex_home).await,
-            memory_root(codex_home).to_path_buf(),
-        )
-    }
-
     /// Returns the rendered developer instruction for read access, if available.
     pub fn read_prompt(&self) -> Option<&str> {
         self.read_prompt.as_deref()
     }
 
-    fn is_read_surface_enabled<C: MemoriesContext>(&self, context: &C) -> bool {
+    // Just for example
+    pub(crate) fn is_read_surface_enabled<C: MemoriesContext>(&self, context: &C) -> bool {
         context.memory_tool_enabled() && context.use_memories()
     }
 }
 
-/// todo Other builder pattern...
 pub async fn extension(codex_home: &AbsolutePathBuf) -> Arc<MemoriesExtension> {
-    Arc::new(MemoriesExtension::from_codex_home(codex_home).await)
+    Arc::new(MemoriesExtension::new(
+        build_memory_tool_developer_instructions(codex_home).await,
+        memory_root(codex_home).to_path_buf(),
+    ))
 }

codex-rs/ext/memories/src/list_tool.rs

@@ -1,3 +1,5 @@
+// This file must be replaced by tools from `codex-memories-mcp` once the extracted tools land. This is just a vibe-coded copy/paster for now.
+
 use std::borrow::Cow;
 use std::future::Future;
 use std::path::Component;

codex-rs/ext/multi-agent-v2/BUILD.bazel

@@ -1,6 +0,0 @@
-load("//:defs.bzl", "codex_rust_crate")
-
-codex_rust_crate(
-    name = "multi-agent-v2",
-    crate_name = "codex_multi_agent_v2",
-)

codex-rs/ext/multi-agent-v2/Cargo.toml

@@ -1,15 +0,0 @@
-[package]
-edition.workspace = true
-license.workspace = true
-name = "codex-multi-agent-v2"
-version.workspace = true
-
-[lib]
-name = "codex_multi_agent_v2"
-path = "src/lib.rs"
-
-[lints]
-workspace = true
-
-[dependencies]
-codex-extension-api = { workspace = true }

codex-rs/ext/multi-agent-v2/src/lib.rs

@@ -1,74 +0,0 @@
-//! MultiAgentV2 usage-hint prompt contribution packaged as a Codex extension.
-
-#![forbid(unsafe_code)]
-
-use std::sync::Arc;
-
-use codex_extension_api::CodexExtension;
-use codex_extension_api::ContextContributor;
-use codex_extension_api::ExtensionRegistryBuilder;
-use codex_extension_api::PromptFragment;
-
-/// Which kind of agent should receive a MultiAgentV2 usage hint.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub enum UsageHintAudience {
-    RootAgent,
-    SpawnedSubagent,
-    Omitted,
-}
-
-/// Runtime facts needed to render MultiAgentV2 usage-hint contributions.
-///
-/// Hosts should return the current effective settings for the thread being
-/// assembled. The extension owns the routing policy that chooses which hint, if
-/// any, is visible for that invocation.
-pub trait MultiAgentV2Context {
-    fn multi_agent_v2_enabled(&self) -> bool;
-    fn usage_hint_audience(&self) -> UsageHintAudience;
-    fn root_agent_usage_hint_text(&self) -> Option<&str>;
-    fn subagent_usage_hint_text(&self) -> Option<&str>;
-}
-
-/// Prompt-only extension that contributes one MultiAgentV2 usage hint when enabled.
-#[derive(Clone, Copy, Debug, Default)]
-pub struct MultiAgentV2Extension;
-
-impl MultiAgentV2Extension {
-    /// Creates an extension instance.
-    pub fn new() -> Self {
-        Self
-    }
-
-    /// Returns the usage hint text that should be shown for this agent, if any.
-    pub fn usage_hint_text<'a, C: MultiAgentV2Context>(&self, context: &'a C) -> Option<&'a str> {
-        if !context.multi_agent_v2_enabled() {
-            return None;
-        }
-
-        match context.usage_hint_audience() {
-            UsageHintAudience::RootAgent => context.root_agent_usage_hint_text(),
-            UsageHintAudience::SpawnedSubagent => context.subagent_usage_hint_text(),
-            UsageHintAudience::Omitted => None,
-        }
-    }
-}
-
-impl<C: MultiAgentV2Context> ContextContributor<C> for MultiAgentV2Extension {
-    fn contribute(&self, context: &C) -> Vec<PromptFragment> {
-        self.usage_hint_text(context)
-            .map(PromptFragment::separate_developer)
-            .into_iter()
-            .collect()
-    }
-}
-
-impl<C: MultiAgentV2Context> CodexExtension<C> for MultiAgentV2Extension {
-    fn install(self: Arc<Self>, registry: &mut ExtensionRegistryBuilder<C>) {
-        registry.prompt_contributor(self);
-    }
-}
-
-/// Creates a shared MultiAgentV2 extension instance.
-pub fn extension() -> Arc<MultiAgentV2Extension> {
-    Arc::new(MultiAgentV2Extension::new())
-}