Codex Handbook
ext/memories/src/extension.rs 128 lines
use std::sync::Arc;use codex_core::config::Config;use codex_extension_api::ConfigContributor;use codex_extension_api::ContextContributor;use codex_extension_api::ExtensionData;use codex_extension_api::ExtensionFuture;use codex_extension_api::ExtensionRegistryBuilder;use codex_extension_api::PromptFragment;use codex_extension_api::ThreadLifecycleContributor;use codex_extension_api::ThreadStartInput;use codex_extension_api::ToolContributor;use codex_features::Feature;use codex_otel::MetricsClient;use codex_utils_absolute_path::AbsolutePathBuf;use crate::local::LocalMemoriesBackend;use crate::prompts::build_memory_tool_developer_instructions;use crate::tools;/// Contributes Codex memory read-path prompt context and memory read tools.#[derive(Clone, Default)]pub(crate) struct MemoriesExtension {    metrics_client: Option<MetricsClient>,}impl MemoriesExtension {    fn new(metrics_client: Option<MetricsClient>) -> Self {        Self { metrics_client }    }}#[derive(Clone, Debug)]pub(crate) struct MemoriesExtensionConfig {    pub(crate) enabled: bool,    pub(crate) dedicated_tools: bool,    pub(crate) codex_home: AbsolutePathBuf,}impl MemoriesExtensionConfig {    fn from_config(config: &Config) -> Self {        Self {            enabled: config.features.enabled(Feature::MemoryTool) && config.memories.use_memories,            dedicated_tools: config.memories.dedicated_tools,            codex_home: config.codex_home.clone(),        }    }}impl ContextContributor for MemoriesExtension {    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 {            let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {                return Vec::new();            };            if !config.enabled {                return Vec::new();            }            build_memory_tool_developer_instructions(&config.codex_home)                .await                .map(PromptFragment::developer_policy)                .into_iter()                .collect()        })    }}impl ThreadLifecycleContributor<Config> for MemoriesExtension {    fn on_thread_start<'a>(        &'a self,        input: ThreadStartInput<'a, Config>,    ) -> ExtensionFuture<'a, ()> {        Box::pin(async move {            input                .thread_store                .insert(MemoriesExtensionConfig::from_config(input.config));        })    }}impl ConfigContributor<Config> for MemoriesExtension {    fn on_config_changed(        &self,        _session_store: &ExtensionData,        thread_store: &ExtensionData,        _previous_config: &Config,        new_config: &Config,    ) {        thread_store.insert(MemoriesExtensionConfig::from_config(new_config));    }}impl ToolContributor for MemoriesExtension {    fn tools(        &self,        _session_store: &ExtensionData,        thread_store: &ExtensionData,    ) -> Vec<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {        let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {            return Vec::new();        };        if !config.enabled || !config.dedicated_tools {            return Vec::new();        }        tools::memory_tools(            LocalMemoriesBackend::from_codex_home(&config.codex_home),            self.metrics_client.clone(),        )    }}/// Installs the memories extension contributors into the extension registry.pub fn install(    registry: &mut ExtensionRegistryBuilder<Config>,    metrics_client: Option<MetricsClient>,) {    let extension = Arc::new(MemoriesExtension::new(metrics_client));    registry.thread_lifecycle_contributor(extension.clone());    registry.config_contributor(extension.clone());    registry.prompt_contributor(extension.clone());    registry.tool_contributor(extension);}