Fix memories extension tests

codex-rs/config/src/schema.rs

@@ -34,6 +34,15 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema {
             );
             continue;
         }
+        if feature.id == codex_features::Feature::MemoryTool {
+            validation.properties.insert(
+                feature.key.to_string(),
+                schema_gen.subschema_for::<codex_features::FeatureToml<
+                    codex_features::MemoriesFeatureConfigToml,
+                >>(),
+            );
+            continue;
+        }
         if feature.id == codex_features::Feature::AppsMcpPathOverride {
             validation.properties.insert(
                 feature.key.to_string(),

codex-rs/core/config.schema.json

@@ -483,7 +483,7 @@
               "type": "boolean"
             },
             "memories": {
-              "type": "boolean"
+              "$ref": "#/definitions/FeatureToml_for_MemoriesFeatureConfigToml"
             },
             "memory_tool": {
               "type": "boolean"
@@ -797,6 +797,16 @@
         }
       ]
     },
+    "FeatureToml_for_MemoriesFeatureConfigToml": {
+      "anyOf": [
+        {
+          "type": "boolean"
+        },
+        {
+          "$ref": "#/definitions/MemoriesFeatureConfigToml"
+        }
+      ]
+    },
     "FeatureToml_for_MultiAgentV2ConfigToml": {
       "anyOf": [
         {
@@ -1272,6 +1282,18 @@
       },
       "type": "object"
     },
+    "MemoriesFeatureConfigToml": {
+      "additionalProperties": false,
+      "properties": {
+        "custom_tools": {
+          "type": "boolean"
+        },
+        "enabled": {
+          "type": "boolean"
+        }
+      },
+      "type": "object"
+    },
     "MemoriesToml": {
       "additionalProperties": false,
       "description": "Memories settings loaded from config.toml.",
@@ -4406,7 +4428,7 @@
           "type": "boolean"
         },
         "memories": {
-          "type": "boolean"
+          "$ref": "#/definitions/FeatureToml_for_MemoriesFeatureConfigToml"
         },
         "memory_tool": {
           "type": "boolean"

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

@@ -32,13 +32,16 @@ impl MemoriesExtension {
 #[derive(Clone, Debug)]
 pub(crate) struct MemoriesExtensionConfig {
     pub(crate) enabled: bool,
+    pub(crate) custom_tools_enabled: bool,
     pub(crate) codex_home: AbsolutePathBuf,
 }
 
 impl MemoriesExtensionConfig {
     fn from_config(config: &Config) -> Self {
+        let enabled = config.features.enabled(Feature::MemoryTool) && config.memories.use_memories;
         Self {
-            enabled: config.features.enabled(Feature::MemoryTool) && config.memories.use_memories,
+            enabled,
+            custom_tools_enabled: enabled && config.features.memories_custom_tools(),
             codex_home: config.codex_home.clone(),
         }
     }
@@ -97,7 +100,7 @@ impl ToolContributor for MemoriesExtension {
         let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
             return Vec::new();
         };
-        if !config.enabled {
+        if !config.custom_tools_enabled {
             return Vec::new();
         }
 
@@ -116,7 +119,6 @@ pub fn install(
     let extension = Arc::new(MemoriesExtension::new(metrics_client));
     registry.thread_lifecycle_contributor(extension.clone());
     registry.config_contributor(extension.clone());
-    registry.prompt_contributor(extension);
-    // Keep the read/retrieval tools out of app-server until that rollout is intentional.
-    // registry.tool_contributor(extension);
+    registry.prompt_contributor(extension.clone());
+    registry.tool_contributor(extension);
 }

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

@@ -3,6 +3,7 @@ use std::sync::Arc;
 
 use codex_extension_api::ContextContributor;
 use codex_extension_api::ExtensionData;
+use codex_extension_api::ExtensionRegistryBuilder;
 use codex_extension_api::PromptSlot;
 use codex_extension_api::ToolCall;
 use codex_extension_api::ToolContributor;
@@ -21,6 +22,15 @@ use crate::extension::MemoriesExtension;
 use crate::extension::MemoriesExtensionConfig;
 use crate::local::LocalMemoriesBackend;
 
+#[test]
+fn install_registers_tool_contributor() {
+    let mut builder = ExtensionRegistryBuilder::<codex_core::config::Config>::new();
+    crate::install(&mut builder, None);
+    let registry = builder.build();
+
+    assert_eq!(registry.tool_contributors().len(), 1);
+}
+
 #[test]
 fn tools_are_not_contributed_without_thread_config() {
     let extension = MemoriesExtension::default();
@@ -41,6 +51,7 @@ fn tools_are_not_contributed_when_disabled() {
     let thread_store = ExtensionData::new("thread");
     thread_store.insert(MemoriesExtensionConfig {
         enabled: false,
+        custom_tools_enabled: false,
         codex_home: test_path_buf("/tmp/codex-home").abs(),
     });
 
@@ -57,6 +68,7 @@ fn tools_are_contributed_when_enabled() {
     let thread_store = ExtensionData::new("thread");
     thread_store.insert(MemoriesExtensionConfig {
         enabled: true,
+        custom_tools_enabled: true,
         codex_home: test_path_buf("/tmp/codex-home").abs(),
     });
 
@@ -77,6 +89,23 @@ fn tools_are_contributed_when_enabled() {
     );
 }
 
+#[test]
+fn tools_are_not_contributed_when_custom_tools_disabled() {
+    let extension = MemoriesExtension::default();
+    let thread_store = ExtensionData::new("thread");
+    thread_store.insert(MemoriesExtensionConfig {
+        enabled: true,
+        custom_tools_enabled: false,
+        codex_home: test_path_buf("/tmp/codex-home").abs(),
+    });
+
+    assert!(
+        extension
+            .tools(&ExtensionData::new("session"), &thread_store)
+            .is_empty()
+    );
+}
+
 #[test]
 fn ad_hoc_tool_definition_includes_filename_contract() {
     let tool = memory_tool(
@@ -115,6 +144,7 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() {
     let thread_store = ExtensionData::new("thread");
     thread_store.insert(MemoriesExtensionConfig {
         enabled: true,
+        custom_tools_enabled: false,
         codex_home: tempdir.path().abs(),
     });
 

codex-rs/features/src/feature_configs.rs

@@ -48,6 +48,25 @@ impl FeatureConfig for MultiAgentV2ConfigToml {
     }
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)]
+#[serde(deny_unknown_fields)]
+pub struct MemoriesFeatureConfigToml {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub enabled: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub custom_tools: Option<bool>,
+}
+
+impl FeatureConfig for MemoriesFeatureConfigToml {
+    fn enabled(&self) -> Option<bool> {
+        self.enabled
+    }
+
+    fn set_enabled(&mut self, enabled: bool) {
+        self.enabled = Some(enabled);
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)]
 #[serde(deny_unknown_fields)]
 pub struct AppsMcpPathOverrideConfigToml {

codex-rs/features/src/lib.rs

@@ -17,6 +17,7 @@ use toml::Table;
 mod feature_configs;
 mod legacy;
 pub use feature_configs::AppsMcpPathOverrideConfigToml;
+pub use feature_configs::MemoriesFeatureConfigToml;
 pub use feature_configs::MultiAgentV2ConfigToml;
 pub use feature_configs::NetworkProxyConfigToml;
 pub use feature_configs::NetworkProxyDomainPermissionToml;
@@ -284,6 +285,7 @@ pub struct LegacyFeatureUsage {
 pub struct Features {
     enabled: BTreeSet<Feature>,
     legacy_usages: BTreeSet<LegacyFeatureUsage>,
+    memories_custom_tools: bool,
 }
 
 #[derive(Debug, Clone, Default)]
@@ -322,6 +324,7 @@ impl Features {
         Self {
             enabled: set,
             legacy_usages: BTreeSet::new(),
+            memories_custom_tools: false,
         }
     }
 
@@ -337,6 +340,10 @@ impl Features {
         self.enabled(Feature::UseLegacyLandlock)
     }
 
+    pub fn memories_custom_tools(&self) -> bool {
+        self.memories_custom_tools
+    }
+
     pub fn enable(&mut self, f: Feature) -> &mut Self {
         self.enabled.insert(f);
         self
@@ -586,6 +593,8 @@ pub fn is_known_feature_key(key: &str) -> bool {
 /// Deserializable features table for TOML.
 #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)]
 pub struct FeaturesToml {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub memories: Option<FeatureToml<MemoriesFeatureConfigToml>>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub multi_agent_v2: Option<FeatureToml<MultiAgentV2ConfigToml>>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
@@ -600,12 +609,18 @@ impl Features {
     fn apply_toml(&mut self, features: &FeaturesToml) {
         let entries = features.entries();
         self.apply_map(&entries);
+        if let Some(custom_tools) = features.memories_custom_tools() {
+            self.memories_custom_tools = custom_tools;
+        }
     }
 }
 
 impl FeaturesToml {
     pub fn entries(&self) -> BTreeMap<String, bool> {
         let mut entries = self.entries.clone();
+        if let Some(enabled) = self.memories.as_ref().and_then(FeatureToml::enabled) {
+            entries.insert(Feature::MemoryTool.key().to_string(), enabled);
+        }
         if let Some(enabled) = self.multi_agent_v2.as_ref().and_then(FeatureToml::enabled) {
             entries.insert(Feature::MultiAgentV2.key().to_string(), enabled);
         }
@@ -622,8 +637,16 @@ impl FeaturesToml {
         entries
     }
 
+    pub fn memories_custom_tools(&self) -> Option<bool> {
+        match self.memories.as_ref()? {
+            FeatureToml::Enabled(_) => None,
+            FeatureToml::Config(config) => config.custom_tools,
+        }
+    }
+
     pub fn materialize_resolved_enabled(&mut self, features: &Features) {
         let Self {
+            memories,
             multi_agent_v2,
             apps_mcp_path_override,
             network_proxy,
@@ -634,7 +657,9 @@ impl FeaturesToml {
         }
         for spec in FEATURES {
             let enabled = features.enabled(spec.id);
-            if spec.id == Feature::MultiAgentV2 {
+            if spec.id == Feature::MemoryTool {
+                materialize_resolved_memories_feature(memories, features, enabled);
+            } else if spec.id == Feature::MultiAgentV2 {
                 materialize_resolved_feature_enabled(multi_agent_v2, enabled);
             } else if spec.id == Feature::AppsMcpPathOverride {
                 materialize_resolved_feature_enabled(apps_mcp_path_override, enabled);
@@ -657,6 +682,30 @@ fn materialize_resolved_feature_enabled<T: FeatureConfig>(
     }
 }
 
+fn materialize_resolved_memories_feature(
+    feature: &mut Option<FeatureToml<MemoriesFeatureConfigToml>>,
+    features: &Features,
+    enabled: bool,
+) {
+    if !features.memories_custom_tools {
+        materialize_resolved_feature_enabled(feature, enabled);
+        return;
+    }
+
+    match feature {
+        Some(FeatureToml::Config(config)) => {
+            config.set_enabled(enabled);
+            config.custom_tools = Some(true);
+        }
+        Some(FeatureToml::Enabled(_)) | None => {
+            *feature = Some(FeatureToml::Config(MemoriesFeatureConfigToml {
+                enabled: Some(enabled),
+                custom_tools: Some(true),
+            }));
+        }
+    }
+}
+
 impl From<BTreeMap<String, bool>> for FeaturesToml {
     fn from(entries: BTreeMap<String, bool>) -> Self {
         Self {

codex-rs/features/src/tests.rs

@@ -532,6 +532,56 @@ multi_agent_v2 = true
     assert_eq!(features.multi_agent_v2, Some(FeatureToml::Enabled(true)));
 }
 
+#[test]
+fn memories_feature_config_deserializes_table() {
+    let features: FeaturesToml = toml::from_str(
+        r#"
+[memories]
+enabled = true
+custom_tools = true
+"#,
+    )
+    .expect("features table should deserialize");
+
+    assert_eq!(
+        features.entries(),
+        BTreeMap::from([("memories".to_string(), true)])
+    );
+    assert_eq!(features.memories_custom_tools(), Some(true));
+    assert_eq!(
+        features.memories,
+        Some(crate::FeatureToml::Config(
+            crate::MemoriesFeatureConfigToml {
+                enabled: Some(true),
+                custom_tools: Some(true),
+            }
+        ))
+    );
+}
+
+#[test]
+fn memories_custom_tools_config_does_not_enable_feature() {
+    let features_toml: FeaturesToml = toml::from_str(
+        r#"
+[memories]
+custom_tools = true
+"#,
+    )
+    .expect("features table should deserialize");
+    let features = Features::from_sources(
+        FeatureConfigSource {
+            features: Some(&features_toml),
+            ..Default::default()
+        },
+        FeatureConfigSource::default(),
+        FeatureOverrides::default(),
+    );
+
+    assert!(!features.enabled(Feature::MemoryTool));
+    assert!(features.memories_custom_tools());
+    assert_eq!(features_toml.entries(), BTreeMap::new());
+}
+
 #[test]
 fn multi_agent_v2_feature_config_deserializes_table() {
     let features: FeaturesToml = toml::from_str(