diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 19ef21197d..518b67f673 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1868,11 +1868,13 @@ dependencies = [ "codex-core-plugins", "codex-device-key", "codex-exec-server", + "codex-extension-api", "codex-external-agent-migration", "codex-external-agent-sessions", "codex-features", "codex-feedback", "codex-file-search", + "codex-git-attribution", "codex-git-utils", "codex-hooks", "codex-login", @@ -2451,7 +2453,6 @@ dependencies = [ "codex-hooks", "codex-login", "codex-mcp", - "codex-memories", "codex-memories-read", "codex-model-provider", "codex-model-provider-info", @@ -2550,6 +2551,7 @@ dependencies = [ "codex-config", "codex-core", "codex-exec-server", + "codex-extension-api", "codex-features", "codex-login", "codex-model-provider-info", @@ -2786,10 +2788,7 @@ dependencies = [ name = "codex-extension-api" version = "0.0.0" dependencies = [ - "codex-git-attribution", - "codex-guardian", - "codex-memories", - "codex-multi-agent-v2", + "codex-protocol", "codex-tools", "serde_json", "thiserror 2.0.18", @@ -2877,7 +2876,9 @@ dependencies = [ name = "codex-git-attribution" version = "0.0.0" dependencies = [ + "codex-core", "codex-extension-api", + "codex-features", ] [[package]] @@ -2904,13 +2905,6 @@ dependencies = [ "walkdir", ] -[[package]] -name = "codex-guardian" -version = "0.0.0" -dependencies = [ - "codex-extension-api", -] - [[package]] name = "codex-hooks" version = "0.0.0" @@ -3072,6 +3066,7 @@ dependencies = [ "codex-config", "codex-core", "codex-exec-server", + "codex-extension-api", "codex-login", "codex-protocol", "codex-shell-command", @@ -3098,12 +3093,17 @@ dependencies = [ name = "codex-memories" version = "0.0.0" dependencies = [ + "codex-core", "codex-extension-api", + "codex-features", + "codex-memories-mcp", "codex-memories-read", + "codex-protocol", + "codex-tools", "codex-utils-absolute-path", - "rmcp", + "codex-utils-stream-parser", + "serde", "serde_json", - "thiserror 2.0.18", "tokio", ] @@ -3248,13 +3248,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "codex-multi-agent-v2" -version = "0.0.0" -dependencies = [ - "codex-extension-api", -] - [[package]] name = "codex-network-proxy" version = "0.0.0" @@ -4305,6 +4298,7 @@ dependencies = [ "codex-config", "codex-core", "codex-exec-server", + "codex-extension-api", "codex-features", "codex-hooks", "codex-login", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 5c2572f16f..e3a87c048e 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -46,10 +46,8 @@ members = [ "execpolicy", "execpolicy-legacy", "ext/extension-api", - "ext/guardian", "ext/git-attribution", "ext/memories", - "ext/multi-agent-v2", "external-agent-migration", "external-agent-sessions", "keyring-store", @@ -169,10 +167,8 @@ codex-external-agent-migration = { path = "external-agent-migration" } codex-external-agent-sessions = { path = "external-agent-sessions" } codex-experimental-api-macros = { path = "codex-experimental-api-macros" } codex-features = { path = "features" } -codex-guardian = { path = "ext/guardian" } codex-git-attribution = { path = "ext/git-attribution" } codex-memories = { path = "ext/memories" } -codex-multi-agent-v2 = { path = "ext/multi-agent-v2" } codex-feedback = { path = "feedback" } codex-install-context = { path = "install-context" } codex-file-search = { path = "file-search" } diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index 3d6f64793e..9449f8dcad 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -1,12 +1,12 @@ use std::sync::Arc; -use codex_core::ExtensionContext; +use codex_core::config::Config; use codex_extension_api::ExtensionRegistry; use codex_extension_api::ExtensionRegistryBuilder; -pub(crate) fn thread_extensions() -> Arc> { +pub(crate) fn thread_extensions() -> Arc> { Arc::new( - ExtensionRegistryBuilder::::new() + ExtensionRegistryBuilder::::new() .with_extension(codex_git_attribution::extension()) .build(), ) diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 6a28c31e2b..46a75111be 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -99,6 +99,7 @@ async fn queue_refresh( #[cfg(test)] mod tests { use super::*; + use crate::extensions::thread_extensions; use async_trait::async_trait; use codex_arg0::Arg0DispatchPaths; use codex_config::CloudRequirementsLoader; @@ -120,7 +121,6 @@ mod tests { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use tempfile::TempDir; - use crate::extensions::thread_extensions; #[tokio::test] async fn strict_refresh_reports_thread_planning_failures() -> anyhow::Result<()> { diff --git a/codex-rs/core-api/Cargo.toml b/codex-rs/core-api/Cargo.toml index 0cc084650c..1889daf514 100644 --- a/codex-rs/core-api/Cargo.toml +++ b/codex-rs/core-api/Cargo.toml @@ -18,6 +18,7 @@ codex-arg0 = { workspace = true } codex-analytics = { workspace = true } codex-config = { workspace = true } codex-core = { workspace = true } +codex-extension-api = { workspace = true } codex-exec-server = { workspace = true } codex-features = { workspace = true } codex-login = { workspace = true } diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index f9bdc9b56b..0dccf1a56a 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -46,6 +46,7 @@ pub use codex_core::thread_store_from_config; pub use codex_exec_server::EnvironmentManager; pub use codex_exec_server::EnvironmentManagerArgs; pub use codex_exec_server::ExecServerRuntimePaths; +pub use codex_extension_api::empty_extension_registry; pub use codex_features::Feature; pub use codex_features::Features; pub use codex_login::AuthManager; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 44c6aacac5..7b03d3045e 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -36,6 +36,7 @@ codex-config = { workspace = true } codex-core-plugins = { workspace = true } codex-core-skills = { workspace = true } codex-exec-server = { workspace = true } +codex-extension-api = { workspace = true } codex-features = { workspace = true } codex-feedback = { workspace = true } codex-login = { workspace = true } diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index a89d8fc973..039c0de9d2 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -83,6 +83,7 @@ pub(crate) async fn run_codex_thread_interactive( skills_manager: Arc::clone(&parent_session.services.skills_manager), plugins_manager: Arc::clone(&parent_session.services.plugins_manager), mcp_manager: Arc::clone(&parent_session.services.mcp_manager), + extensions: Arc::clone(&parent_session.services.extensions), skills_watcher: Arc::clone(&parent_session.services.skills_watcher), conversation_history: initial_history.unwrap_or(InitialHistory::New), session_source: SessionSource::SubAgent(subagent_source.clone()), diff --git a/codex-rs/core/src/commit_attribution.rs b/codex-rs/core/src/commit_attribution.rs deleted file mode 100644 index 8df9661431..0000000000 --- a/codex-rs/core/src/commit_attribution.rs +++ /dev/null @@ -1,33 +0,0 @@ -const DEFAULT_ATTRIBUTION_VALUE: &str = "Codex "; - -fn build_commit_message_trailer(config_attribution: Option<&str>) -> Option { - let value = resolve_attribution_value(config_attribution)?; - Some(format!("Co-authored-by: {value}")) -} - -pub(crate) fn commit_message_trailer_instruction( - config_attribution: Option<&str>, -) -> Option { - let trailer = build_commit_message_trailer(config_attribution)?; - Some(format!( - "When you write or edit a git commit message, ensure the message ends with this trailer exactly once:\n{trailer}\n\nRules:\n- Keep existing trailers and append this trailer at the end if missing.\n- Do not duplicate this trailer if it already exists.\n- Keep one blank line between the commit body and trailer block." - )) -} - -fn resolve_attribution_value(config_attribution: Option<&str>) -> Option { - match config_attribution { - Some(value) => { - let trimmed = value.trim(); - if trimmed.is_empty() { - None - } else { - Some(trimmed.to_string()) - } - } - None => Some(DEFAULT_ATTRIBUTION_VALUE.to_string()), - } -} - -#[cfg(test)] -#[path = "commit_attribution_tests.rs"] -mod tests; diff --git a/codex-rs/core/src/commit_attribution_tests.rs b/codex-rs/core/src/commit_attribution_tests.rs deleted file mode 100644 index 8b2ce27e1f..0000000000 --- a/codex-rs/core/src/commit_attribution_tests.rs +++ /dev/null @@ -1,43 +0,0 @@ -use super::build_commit_message_trailer; -use super::commit_message_trailer_instruction; -use super::resolve_attribution_value; - -#[test] -fn blank_attribution_disables_trailer_prompt() { - assert_eq!(build_commit_message_trailer(Some("")), None); - assert_eq!(commit_message_trailer_instruction(Some(" ")), None); -} - -#[test] -fn default_attribution_uses_codex_trailer() { - assert_eq!( - build_commit_message_trailer(/*config_attribution*/ None).as_deref(), - Some("Co-authored-by: Codex ") - ); -} - -#[test] -fn resolve_value_handles_default_custom_and_blank() { - assert_eq!( - resolve_attribution_value(/*config_attribution*/ None), - Some("Codex ".to_string()) - ); - assert_eq!( - resolve_attribution_value(Some("MyAgent ")), - Some("MyAgent ".to_string()) - ); - assert_eq!( - resolve_attribution_value(Some("MyAgent")), - Some("MyAgent".to_string()) - ); - assert_eq!(resolve_attribution_value(Some(" ")), None); -} - -#[test] -fn instruction_mentions_trailer_and_omits_generated_with() { - let instruction = commit_message_trailer_instruction(Some("AgentX ")) - .expect("instruction expected"); - assert!(instruction.contains("Co-authored-by: AgentX ")); - assert!(instruction.contains("exactly once")); - assert!(!instruction.contains("Generated-with")); -} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 0cdf0e2d46..222e8105e5 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -22,10 +22,10 @@ mod config_lock; pub use codex_thread::CodexThread; pub use codex_thread::CodexThreadTurnContextOverrides; pub use codex_thread::ThreadConfigSnapshot; +pub use session::turn_context::TurnContext; mod agent; mod codex_delegate; mod command_canonicalization; -mod commit_attribution; pub mod config; pub mod connectors; pub mod context; diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 8717427afe..a88d94a1de 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -20,6 +20,7 @@ use crate::session::turn::built_tools; use crate::state_db_bridge::StateDbHandle; use crate::thread_manager::ThreadManager; use crate::thread_manager::thread_store_from_config; +use codex_extension_api::empty_extension_registry; /// Build the model-visible `input` list for a single debug turn. #[doc(hidden)] @@ -45,6 +46,7 @@ pub async fn build_prompt_input( Arc::clone(&auth_manager), SessionSource::Exec, Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store, state_db.clone(), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 1f40438cd6..b03c17b8f7 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -15,7 +15,6 @@ use crate::agent::MailboxReceiver; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; use crate::build_available_skills; -use crate::commit_attribution::commit_message_trailer_instruction; use crate::compact; use crate::config::ManagedFeatures; use crate::config::resolve_tool_suggest_config_from_layer_stack; @@ -53,6 +52,7 @@ use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_exec_server::FileSystemSandboxContext; +use codex_extension_api::PromptSlot; use codex_features::FEATURES; use codex_features::Feature; use codex_features::unstable_features_warning_event; @@ -392,6 +392,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) skills_manager: Arc, pub(crate) plugins_manager: Arc, pub(crate) mcp_manager: Arc, + pub(crate) extensions: Arc>, pub(crate) skills_watcher: Arc, pub(crate) conversation_history: InitialHistory, pub(crate) session_source: SessionSource, @@ -455,6 +456,7 @@ impl Codex { skills_manager, plugins_manager, mcp_manager, + extensions, skills_watcher, conversation_history, session_source, @@ -650,6 +652,7 @@ impl Codex { skills_manager, plugins_manager, mcp_manager.clone(), + extensions, skills_watcher, agent_control, environment_manager, @@ -2556,6 +2559,7 @@ impl Session { ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::with_capacity(2); + let mut separate_developer_sections = Vec::::new(); let ( reference_context_item, previous_turn_settings, @@ -2693,12 +2697,23 @@ impl Session { { developer_sections.push(plugin_instructions.render()); } - if turn_context.features.enabled(Feature::CodexGitCommit) - && let Some(commit_message_instruction) = commit_message_trailer_instruction( - turn_context.config.commit_attribution.as_deref(), - ) - { - developer_sections.push(commit_message_instruction); + for contributor in self.services.extensions.context_contributors() { + for fragment in contributor.contribute( + &self.services.session_extension_data, + &self.services.thread_extension_data, + ) { + match fragment.slot() { + PromptSlot::DeveloperPolicy | PromptSlot::DeveloperCapabilities => { + developer_sections.push(fragment.text().to_string()); + } + PromptSlot::ContextualUser => { + contextual_user_sections.push(fragment.text().to_string()); + } + PromptSlot::SeparateDeveloper => { + separate_developer_sections.push(fragment.text().to_string()); + } + } + } } if let Some(user_instructions) = turn_context.user_instructions.as_deref() { contextual_user_sections.push( @@ -2732,6 +2747,13 @@ impl Session { { items.push(developer_message); } + for section in separate_developer_sections { + if let Some(developer_message) = + crate::context_manager::updates::build_developer_update_item(vec![section]) + { + items.push(developer_message); + } + } if let Some(usage_hint_text) = multi_agent_v2_usage_hint_text && let Some(usage_hint_message) = crate::context_manager::updates::build_developer_update_item(vec![ diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 3c23c6a4e3..0a16c90067 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -364,7 +364,7 @@ impl Session { skills_manager: Arc, plugins_manager: Arc, mcp_manager: Arc, - extensions: Arc>, + extensions: Arc>, skills_watcher: Arc, agent_control: AgentControl, environment_manager: Arc, @@ -811,6 +811,16 @@ impl Session { SessionId::from(thread_id) }; let agent_control = agent_control.with_session_id(session_id); + let session_extension_data = codex_extension_api::ExtensionData::new(); + let thread_extension_data = codex_extension_api::ExtensionData::new(); + for contributor in extensions.thread_start_contributors() { + contributor.contribute( + config.as_ref(), + &session_extension_data, + &thread_extension_data, + ); + } + let services = SessionServices { // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via @@ -847,8 +857,9 @@ impl Session { plugins_manager: Arc::clone(&plugins_manager), mcp_manager: Arc::clone(&mcp_manager), extensions, - session_extension_data: codex_extension_api::ExtensionData::new(), - thread_extension_data: codex_extension_api::ExtensionData::new(), + // TODO(jif): extract session to share between sub-agents + session_extension_data, + thread_extension_data, skills_watcher, agent_control, network_proxy, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 846489647c..587a21fbea 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6003,18 +6003,43 @@ async fn make_multi_agent_v2_usage_hint_test_session( (session, turn_context) } +struct GitAttributionTestContributor; +struct GitAttributionTestState; + +impl codex_extension_api::ContextContributor for GitAttributionTestContributor { + fn contribute( + &self, + _session_store: &codex_extension_api::ExtensionData, + thread_store: &codex_extension_api::ExtensionData, + ) -> Vec { + thread_store + .get::() + .is_some() + .then(|| { + codex_extension_api::PromptFragment::developer_policy( + "git attribution extension enabled", + ) + }) + .into_iter() + .collect() + } +} + +fn git_attribution_test_registry() +-> Arc> { + let mut builder = codex_extension_api::ExtensionRegistryBuilder::new(); + builder.prompt_contributor(Arc::new(GitAttributionTestContributor)); + Arc::new(builder.build()) +} + #[tokio::test] async fn build_initial_context_includes_git_attribution_from_extensions() { - let (mut session, mut turn_context) = make_session_and_context().await; - session.services.extensions = Arc::new( - codex_extension_api::ExtensionRegistryBuilder::::new() - .with_extension(codex_git_attribution::extension()) - .build(), - ); - turn_context - .features - .enable(Feature::CodexGitCommit) - .expect("test features should allow git attribution"); + let (mut session, turn_context) = make_session_and_context().await; + session.services.extensions = git_attribution_test_registry(); + session + .services + .thread_extension_data + .insert(GitAttributionTestState); let initial_context = session.build_initial_context(&turn_context).await; let developer_messages = developer_message_texts(&initial_context); @@ -6023,7 +6048,7 @@ async fn build_initial_context_includes_git_attribution_from_extensions() { developer_messages .iter() .flatten() - .any(|text| text.contains("Co-authored-by: Codex ")), + .any(|text| *text == "git attribution extension enabled"), "expected git attribution developer text, got {developer_messages:?}" ); } @@ -6031,11 +6056,7 @@ async fn build_initial_context_includes_git_attribution_from_extensions() { #[tokio::test] async fn build_initial_context_omits_git_attribution_when_feature_is_disabled() { let (mut session, turn_context) = make_session_and_context().await; - session.services.extensions = Arc::new( - codex_extension_api::ExtensionRegistryBuilder::::new() - .with_extension(codex_git_attribution::extension()) - .build(), - ); + session.services.extensions = git_attribution_test_registry(); let initial_context = session.build_initial_context(&turn_context).await; let developer_messages = developer_message_texts(&initial_context); @@ -6044,7 +6065,7 @@ async fn build_initial_context_omits_git_attribution_when_feature_is_disabled() !developer_messages .iter() .flatten() - .any(|text| text.contains("Co-authored-by:")), + .any(|text| *text == "git attribution extension enabled"), "did not expect git attribution developer text, got {developer_messages:?}" ); } diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 66b85b432f..6ae3ef0e13 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -743,7 +743,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { skills_manager, plugins_manager, mcp_manager, - extensions: crate::empty_extension_registry(), + extensions: codex_extension_api::empty_extension_registry(), skills_watcher, conversation_history: InitialHistory::New, session_source: SessionSource::SubAgent(SubAgentSource::Other( diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index d4fe30063f..0e271671ac 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -52,11 +52,11 @@ impl TurnEnvironment { /// The context needed for a single turn of the thread. #[derive(Debug)] -pub(crate) struct TurnContext { +pub struct TurnContext { pub(crate) sub_id: String, pub(crate) trace_id: Option, pub(crate) realtime_active: bool, - pub(crate) config: Arc, + pub config: Arc, pub(crate) auth_manager: Option>, pub(crate) model_info: ModelInfo, pub(crate) session_telemetry: SessionTelemetry, @@ -84,7 +84,7 @@ pub(crate) struct TurnContext { pub(crate) windows_sandbox_level: WindowsSandboxLevel, pub(crate) shell_environment_policy: ShellEnvironmentPolicy, pub(crate) tools_config: ToolsConfig, - pub(crate) features: ManagedFeatures, + pub features: ManagedFeatures, pub(crate) ghost_snapshot: GhostSnapshotConfig, pub(crate) final_output_json_schema: Option, pub(crate) codex_self_exe: Option, diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index d89f74c015..a255ee9603 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -6,7 +6,6 @@ use crate::agent::AgentControl; use crate::client::ModelClient; use crate::config::StartedNetworkProxy; use crate::exec_policy::ExecPolicyManager; -use crate::extensions::ExtensionContext; use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; @@ -62,7 +61,7 @@ pub(crate) struct SessionServices { pub(crate) skills_manager: Arc, pub(crate) plugins_manager: Arc, pub(crate) mcp_manager: Arc, - pub(crate) extensions: Arc>, + pub(crate) extensions: Arc>, pub(crate) session_extension_data: ExtensionData, pub(crate) thread_extension_data: ExtensionData, pub(crate) skills_watcher: Arc, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index feda91ce2f..a214170bd5 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -23,6 +23,8 @@ use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; +use codex_extension_api::ExtensionRegistry; +use codex_extension_api::empty_extension_registry; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider::create_model_provider; @@ -246,6 +248,7 @@ pub(crate) struct ThreadManagerState { skills_manager: Arc, plugins_manager: Arc, mcp_manager: Arc, + extensions: Arc>, skills_watcher: Arc, thread_store: Arc, session_source: SessionSource, @@ -288,6 +291,7 @@ impl ThreadManager { auth_manager: Arc, session_source: SessionSource, environment_manager: Arc, + extensions: Arc>, analytics_events_client: Option, thread_store: Arc, state_db: Option, @@ -316,6 +320,7 @@ impl ThreadManager { skills_manager, plugins_manager, mcp_manager, + extensions, skills_watcher, thread_store, auth_manager, @@ -417,6 +422,7 @@ impl ThreadManager { skills_manager, plugins_manager, mcp_manager, + extensions: empty_extension_registry(), skills_watcher, thread_store, auth_manager, @@ -1171,6 +1177,7 @@ impl ThreadManagerState { skills_manager: Arc::clone(&self.skills_manager), plugins_manager: Arc::clone(&self.plugins_manager), mcp_manager: Arc::clone(&self.mcp_manager), + extensions: Arc::clone(&self.extensions), skills_watcher: Arc::clone(&self.skills_watcher), conversation_history: initial_history, session_source, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 17ac81a93b..88e21941bc 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -7,6 +7,7 @@ use crate::session::session::SessionSettingsUpdate; use crate::session::tests::make_session_and_context; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; +use codex_extension_api::empty_extension_registry; use codex_features::Feature; use codex_models_manager::manager::RefreshStrategy; use codex_protocol::models::ContentItem; @@ -396,6 +397,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, @@ -512,6 +514,7 @@ async fn explicit_installation_id_skips_codex_home_file() { auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store, state_db.clone(), @@ -549,6 +552,7 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, @@ -604,6 +608,7 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, @@ -666,6 +671,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store, state_db.clone(), @@ -750,6 +756,7 @@ async fn new_uses_active_provider_for_model_refresh() { auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, @@ -964,6 +971,7 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, state_db.clone()), state_db.clone(), @@ -1070,6 +1078,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, state_db.clone()), state_db.clone(), @@ -1165,6 +1174,7 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, state_db.clone()), state_db.clone(), @@ -1306,6 +1316,7 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { auth_manager.clone(), SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, state_db.clone()), state_db.clone(), diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index b80064f1ab..01ff66a752 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -2,7 +2,6 @@ use super::*; use crate::ThreadManager; use crate::config::AgentRoleConfig; use crate::config::DEFAULT_AGENT_MAX_DEPTH; -use crate::empty_extension_registry; use crate::function_tool::FunctionCallError; use crate::init_state_db; use crate::session::tests::make_session_and_context; @@ -16,6 +15,7 @@ use crate::tools::handlers::multi_agents_v2::SendMessageHandler as SendMessageHa use crate::tools::handlers::multi_agents_v2::SpawnAgentHandler as SpawnAgentHandlerV2; use crate::tools::handlers::multi_agents_v2::WaitAgentHandler as WaitAgentHandlerV2; use crate::turn_diff_tracker::TurnDiffTracker; +use codex_extension_api::empty_extension_registry; use codex_features::Feature; use codex_login::AuthManager; use codex_login::CodexAuth; diff --git a/codex-rs/core/tests/common/Cargo.toml b/codex-rs/core/tests/common/Cargo.toml index 8c1a79f02d..50dc2f0ac8 100644 --- a/codex-rs/core/tests/common/Cargo.toml +++ b/codex-rs/core/tests/common/Cargo.toml @@ -17,6 +17,7 @@ base64 = { workspace = true } codex-arg0 = { workspace = true } codex-config = { workspace = true } codex-core = { workspace = true } +codex-extension-api = { workspace = true } codex-exec-server = { workspace = true } codex-features = { workspace = true } codex-hooks = { workspace = true } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 7082d26089..4bdbb5b0b0 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -23,6 +23,7 @@ use codex_core::thread_store_from_config; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::RemoveOptions; +use codex_extension_api::empty_extension_registry; use codex_features::Feature; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; @@ -432,6 +433,7 @@ impl TestCodexBuilder { codex_core::test_support::auth_manager_from_auth(auth.clone()), SessionSource::Exec, Arc::clone(&environment_manager), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store, state_db.clone(), diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 432b57de9f..63b8b3e31a 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -7,6 +7,7 @@ use codex_core::ResponseEvent; use codex_core::ThreadManager; use codex_core::resolve_installation_id; use codex_core::thread_store_from_config; +use codex_extension_api::empty_extension_registry; use codex_features::Feature; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -1122,6 +1123,7 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { auth_manager, SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions.rs b/codex-rs/ext/extension-api/examples/enabled_extensions.rs index cdd4ef5699..f091aa7ab5 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions.rs @@ -13,7 +13,7 @@ fn main() { // 1. Build the extension value owned by the host. let extension = Arc::new(SharedStateExtension); - // 2. Install it into the registry for the context type this host exposes. + // 2. Install it into the registry for the thread-start input type this host exposes. let registry = ExtensionRegistryBuilder::<()>::new() .with_extension(extension) .build(); @@ -63,6 +63,6 @@ fn contribute_prompt( registry .context_contributors() .iter() - .flat_map(|contributor| contributor.contribute(&(), session_store, thread_store)) + .flat_map(|contributor| contributor.contribute(session_store, thread_store)) .collect() } diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs b/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs index e13dfbe8db..9b6612e3ae 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs @@ -22,10 +22,9 @@ impl CodexExtension<()> for SharedStateExtension { #[derive(Debug)] struct StyleContributor; -impl ContextContributor<()> for StyleContributor { +impl ContextContributor for StyleContributor { fn contribute( &self, - _context: &(), session_store: &ExtensionData, thread_store: &ExtensionData, ) -> Vec { @@ -41,10 +40,9 @@ impl ContextContributor<()> for StyleContributor { #[derive(Debug)] struct UsageContributor; -impl ContextContributor<()> for UsageContributor { +impl ContextContributor for UsageContributor { fn contribute( &self, - _context: &(), session_store: &ExtensionData, thread_store: &ExtensionData, ) -> Vec { diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 8e599c805b..34096c9baf 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -13,20 +13,25 @@ pub use tool::ToolCallError; pub use tool::ToolContribution; pub use tool::ToolHandler; +/// Contributor that receives host-owned thread-start input before later +/// contributors read from extension stores. +pub trait ThreadStartContributor: Send + Sync { + fn contribute(&self, input: &C, session_store: &ExtensionData, thread_store: &ExtensionData); +} + /// Extension contribution that adds prompt fragments during prompt assembly. -pub trait ContextContributor: Send + Sync { +pub trait ContextContributor: Send + Sync { fn contribute( &self, - context: &C, session_store: &ExtensionData, thread_store: &ExtensionData, ) -> Vec; } /// Extension contribution that exposes native tools owned by a feature. -pub trait ToolContributor: Send + Sync { +pub trait ToolContributor: Send + Sync { /// Returns the native tools visible for the supplied runtime context. - fn tools(&self, context: &C, thread_store: &ExtensionData) -> Vec>; + fn tools(&self, thread_store: &ExtensionData) -> Vec; } /// Future returned by one ordered turn-item contribution. @@ -38,10 +43,9 @@ pub type TurnItemContributionFuture<'a> = /// Implementations may mutate the item before it is emitted and may use the /// explicitly exposed thread- and turn-lifetime stores when they need durable /// extension-private state. -pub trait TurnItemContributor: Send + Sync { +pub trait TurnItemContributor: Send + Sync { fn contribute<'a>( &'a self, - context: &'a C, thread_store: &'a ExtensionData, turn_store: &'a ExtensionData, item: &'a mut TurnItem, @@ -51,11 +55,10 @@ pub trait TurnItemContributor: Send + Sync { // 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: Send + Sync { +pub trait ApprovalInterceptorContributor: Send + Sync { /// Returns whether this contributor should intercept approvals in `context`. fn intercepts_approvals( &self, - context: &C, thread_store: &ExtensionData, turn_store: &ExtensionData, ) -> bool; diff --git a/codex-rs/ext/extension-api/src/contributors/tool.rs b/codex-rs/ext/extension-api/src/contributors/tool.rs index 28e93bd253..de52f8ecc5 100644 --- a/codex-rs/ext/extension-api/src/contributors/tool.rs +++ b/codex-rs/ext/extension-api/src/contributors/tool.rs @@ -8,14 +8,14 @@ use thiserror::Error; // TMP #[derive(Clone)] -pub struct ToolContribution { +pub struct ToolContribution { spec: ResponsesApiTool, - handler: Arc>, + handler: Arc, supports_parallel_tool_calls: bool, } -impl ToolContribution { - pub fn new(spec: ResponsesApiTool, handler: Arc>) -> Self { +impl ToolContribution { + pub fn new(spec: ResponsesApiTool, handler: Arc) -> Self { Self { spec, handler, @@ -37,17 +37,16 @@ impl ToolContribution { self.supports_parallel_tool_calls } - pub fn handler(&self) -> Arc> { + pub fn handler(&self) -> Arc { Arc::clone(&self.handler) } } //////// Just to make it compile //////////////////////////////// -pub trait ToolHandler: Send + Sync { +pub trait ToolHandler: Send + Sync { /// Handles one JSON-encoded invocation for this tool. fn handle<'a>( &'a self, - context: &'a C, arguments: Value, ) -> Pin> + Send + 'a>>; } diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 11824a598c..0b07f9e285 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -7,6 +7,7 @@ pub use contributors::ApprovalInterceptorContributor; pub use contributors::ContextContributor; pub use contributors::PromptFragment; pub use contributors::PromptSlot; +pub use contributors::ThreadStartContributor; pub use contributors::ToolCallError; pub use contributors::ToolContribution; pub use contributors::ToolContributor; @@ -16,4 +17,5 @@ pub use contributors::TurnItemContributor; pub use extension::CodexExtension; pub use registry::ExtensionRegistry; pub use registry::ExtensionRegistryBuilder; +pub use registry::empty_extension_registry; pub use state::ExtensionData; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index 900c1886bd..4a90e2e606 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -3,20 +3,23 @@ use std::sync::Arc; use crate::ApprovalInterceptorContributor; use crate::CodexExtension; use crate::ContextContributor; +use crate::ThreadStartContributor; use crate::ToolContributor; use crate::TurnItemContributor; /// Mutable registry used while extensions install their typed contributions. pub struct ExtensionRegistryBuilder { - context_contributors: Vec>>, - tool_contributors: Vec>>, - turn_item_contributors: Vec>>, - approval_interceptor_contributors: Vec>>, + thread_start_contributors: Vec>>, + context_contributors: Vec>, + tool_contributors: Vec>, + turn_item_contributors: Vec>, + approval_interceptor_contributors: Vec>, } impl Default for ExtensionRegistryBuilder { fn default() -> Self { Self { + thread_start_contributors: Vec::new(), approval_interceptor_contributors: Vec::new(), context_contributors: Vec::new(), tool_contributors: Vec::new(), @@ -52,29 +55,35 @@ impl ExtensionRegistryBuilder { /// Registers one approval interceptor contributor. pub fn approval_interceptor_contributor( &mut self, - contributor: Arc>, + contributor: Arc, ) { self.approval_interceptor_contributors.push(contributor); } + /// Registers one thread-start contributor. + pub fn thread_start_contributor(&mut self, contributor: Arc>) { + self.thread_start_contributors.push(contributor); + } + /// Registers one prompt contributor. - pub fn prompt_contributor(&mut self, contributor: Arc>) { + pub fn prompt_contributor(&mut self, contributor: Arc) { self.context_contributors.push(contributor); } /// Registers one native tool contributor. - pub fn tool_contributor(&mut self, contributor: Arc>) { + pub fn tool_contributor(&mut self, contributor: Arc) { self.tool_contributors.push(contributor); } /// Registers one ordered turn-item contributor. - pub fn turn_item_contributor(&mut self, contributor: Arc>) { + pub fn turn_item_contributor(&mut self, contributor: Arc) { self.turn_item_contributors.push(contributor); } /// Finishes construction and returns the immutable registry. pub fn build(self) -> ExtensionRegistry { ExtensionRegistry { + thread_start_contributors: self.thread_start_contributors, approval_interceptor_contributors: self.approval_interceptor_contributors, context_contributors: self.context_contributors, tool_contributors: self.tool_contributors, @@ -85,32 +94,41 @@ impl ExtensionRegistryBuilder { /// Immutable typed registry produced after extensions are installed. pub struct ExtensionRegistry { - context_contributors: Vec>>, - tool_contributors: Vec>>, - turn_item_contributors: Vec>>, - approval_interceptor_contributors: Vec>>, + thread_start_contributors: Vec>>, + context_contributors: Vec>, + tool_contributors: Vec>, + turn_item_contributors: Vec>, + approval_interceptor_contributors: Vec>, } impl ExtensionRegistry { + /// Returns the registered thread-start contributors. + pub fn thread_start_contributors(&self) -> &[Arc>] { + &self.thread_start_contributors + } + /// Returns the registered approval interceptor contributors. - pub fn approval_interceptor_contributors( - &self, - ) -> &[Arc>] { + pub fn approval_interceptor_contributors(&self) -> &[Arc] { &self.approval_interceptor_contributors } /// Returns the registered prompt contributors. - pub fn context_contributors(&self) -> &[Arc>] { + pub fn context_contributors(&self) -> &[Arc] { &self.context_contributors } /// Returns the registered native tool contributors. - pub fn tool_contributors(&self) -> &[Arc>] { + pub fn tool_contributors(&self) -> &[Arc] { &self.tool_contributors } /// Returns the registered ordered turn-item contributors. - pub fn turn_item_contributors(&self) -> &[Arc>] { + pub fn turn_item_contributors(&self) -> &[Arc] { &self.turn_item_contributors } } + +/// Creates an empty shared registry for hosts that do not install extensions. +pub fn empty_extension_registry() -> Arc> { + Arc::new(ExtensionRegistryBuilder::new().build()) +} diff --git a/codex-rs/ext/git-attribution/Cargo.toml b/codex-rs/ext/git-attribution/Cargo.toml index 46bdbbec0d..e5db5a1834 100644 --- a/codex-rs/ext/git-attribution/Cargo.toml +++ b/codex-rs/ext/git-attribution/Cargo.toml @@ -12,4 +12,6 @@ path = "src/lib.rs" workspace = true [dependencies] +codex-core = { workspace = true } codex-extension-api = { workspace = true } +codex-features = { workspace = true } diff --git a/codex-rs/ext/git-attribution/src/lib.rs b/codex-rs/ext/git-attribution/src/lib.rs index dcdafd3dc4..d9a35ae2c8 100644 --- a/codex-rs/ext/git-attribution/src/lib.rs +++ b/codex-rs/ext/git-attribution/src/lib.rs @@ -1,18 +1,16 @@ use std::sync::Arc; +use codex_core::config::Config; 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::ThreadStartContributor; +use codex_features::Feature; const DEFAULT_ATTRIBUTION_VALUE: &str = "Codex "; -pub trait GitAttributionContext { - fn git_attribution_enabled(&self) -> bool; - fn commit_attribution(&self) -> Option<&str>; -} - /// Prompt-only extension that contributes the configured git-attribution instruction. #[derive(Clone, Copy, Debug, Default)] pub struct GitAttributionExtension; @@ -22,35 +20,50 @@ impl GitAttributionExtension { pub fn new() -> Self { Self } - - /// Returns the model-visible trailer instruction, if attribution is enabled. - pub fn instruction(&self, context: &C) -> Option { - if !context.git_attribution_enabled() { - return None; - } - let trailer = build_commit_message_trailer(context.commit_attribution())?; - Some(format!( - "When you write or edit a git commit message, ensure the message ends with this trailer exactly once:\n{trailer}\n\nRules:\n- Keep existing trailers and append this trailer at the end if missing.\n- Do not duplicate this trailer if it already exists.\n- Keep one blank line between the commit body and trailer block." - )) - } } -impl ContextContributor for GitAttributionExtension { +impl ContextContributor for GitAttributionExtension { fn contribute( &self, - context: &C, _session_store: &ExtensionData, - _thread_store: &ExtensionData, + thread_store: &ExtensionData, ) -> Vec { - self.instruction(context) + let Some(config_store) = thread_store.get::() else { + return Vec::new(); + }; + if !config_store.enabled { + return Vec::new(); + } + build_instruction(config_store.prompt.as_deref()) .map(PromptFragment::developer_policy) .into_iter() .collect() } } -impl CodexExtension for GitAttributionExtension { - fn install(self: Arc, registry: &mut ExtensionRegistryBuilder) { +#[derive(Clone, Debug, Default)] +struct GitAttributionConfig { + enabled: bool, + prompt: Option, +} + +impl ThreadStartContributor for GitAttributionExtension { + fn contribute( + &self, + config: &Config, + _session_store: &ExtensionData, + thread_store: &ExtensionData, + ) { + thread_store.insert(GitAttributionConfig { + enabled: config.features.enabled(Feature::CodexGitCommit), + prompt: config.commit_attribution.clone(), + }); + } +} + +impl CodexExtension for GitAttributionExtension { + fn install(self: Arc, registry: &mut ExtensionRegistryBuilder) { + registry.thread_start_contributor(self.clone()); registry.prompt_contributor(self); } } @@ -65,6 +78,13 @@ fn build_commit_message_trailer(config_attribution: Option<&str>) -> Option) -> Option { + let trailer = build_commit_message_trailer(config_attribution)?; + Some(format!( + "When you write or edit a git commit message, ensure the message ends with this trailer exactly once:\n{trailer}\n\nRules:\n- Keep existing trailers and append this trailer at the end if missing.\n- Do not duplicate this trailer if it already exists.\n- Keep one blank line between the commit body and trailer block." + )) +} + fn resolve_attribution_value(config_attribution: Option<&str>) -> Option { match config_attribution { Some(value) => { @@ -78,6 +98,3 @@ fn resolve_attribution_value(config_attribution: Option<&str>) -> Option None => Some(DEFAULT_ATTRIBUTION_VALUE.to_string()), } } - -#[cfg(test)] -mod tests; diff --git a/codex-rs/ext/git-attribution/src/tests.rs b/codex-rs/ext/git-attribution/src/tests.rs deleted file mode 100644 index 56ec3d717e..0000000000 --- a/codex-rs/ext/git-attribution/src/tests.rs +++ /dev/null @@ -1,81 +0,0 @@ -use super::GitAttributionContext; -use super::GitAttributionExtension; -use super::build_commit_message_trailer; -use super::resolve_attribution_value; - -struct TestContext { - git_attribution_enabled: bool, - commit_attribution: Option<&'static str>, -} - -impl GitAttributionContext for TestContext { - fn git_attribution_enabled(&self) -> bool { - self.git_attribution_enabled - } - - fn commit_attribution(&self) -> Option<&str> { - self.commit_attribution - } -} - -#[test] -fn blank_attribution_disables_trailer_prompt() { - assert_eq!(build_commit_message_trailer(Some("")), None); - assert_eq!( - GitAttributionExtension::new().instruction(&TestContext { - git_attribution_enabled: true, - commit_attribution: Some(" "), - }), - None - ); -} - -#[test] -fn disabled_context_omits_instruction() { - assert_eq!( - GitAttributionExtension::new().instruction(&TestContext { - git_attribution_enabled: false, - commit_attribution: None, - }), - None - ); -} - -#[test] -fn default_attribution_uses_codex_trailer() { - assert_eq!( - build_commit_message_trailer(/*config_attribution*/ None).as_deref(), - Some("Co-authored-by: Codex ") - ); -} - -#[test] -fn resolve_value_handles_default_custom_and_blank() { - assert_eq!( - resolve_attribution_value(/*config_attribution*/ None), - Some("Codex ".to_string()) - ); - assert_eq!( - resolve_attribution_value(Some("MyAgent ")), - Some("MyAgent ".to_string()) - ); - assert_eq!( - resolve_attribution_value(Some("MyAgent")), - Some("MyAgent".to_string()) - ); - assert_eq!(resolve_attribution_value(Some(" ")), None); -} - -#[test] -fn instruction_mentions_trailer_and_omits_generated_with() { - let instruction = GitAttributionExtension::new() - .instruction(&TestContext { - git_attribution_enabled: true, - commit_attribution: Some("AgentX "), - }) - .expect("instruction expected"); - - assert!(instruction.contains("Co-authored-by: AgentX ")); - assert!(instruction.contains("exactly once")); - assert!(!instruction.contains("Generated-with")); -} diff --git a/codex-rs/ext/memories/Cargo.toml b/codex-rs/ext/memories/Cargo.toml index 3070066ba6..28413ac3dc 100644 --- a/codex-rs/ext/memories/Cargo.toml +++ b/codex-rs/ext/memories/Cargo.toml @@ -12,7 +12,10 @@ path = "src/lib.rs" workspace = true [dependencies] +codex-core = { workspace = true } codex-extension-api = { workspace = true } +codex-features = { workspace = true } +codex-memories-mcp = { workspace = true } codex-memories-read = { workspace = true } codex-protocol = { workspace = true } codex-tools = { workspace = true } @@ -20,6 +23,4 @@ codex-utils-absolute-path = { workspace = true } codex-utils-stream-parser = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } -thiserror = { workspace = true } tokio = { workspace = true, features = ["fs"] } -tracing = { workspace = true } diff --git a/codex-rs/ext/memories/src/citation_output.rs b/codex-rs/ext/memories/src/citation_output.rs deleted file mode 100644 index b02b11dd63..0000000000 --- a/codex-rs/ext/memories/src/citation_output.rs +++ /dev/null @@ -1,70 +0,0 @@ -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; -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 OutputContributor 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 { - let combined = agent_message - .content - .iter() - .map(|entry| match entry { - AgentMessageContent::Text { text } => text.as_str(), - }) - .collect::(); - 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 { - let turn_state = stores.get_or_init::(Default::default); - if turn_state.had_citation.swap(true, Ordering::Relaxed) { - return None; - } - - let thread_stats = stores.get_or_init::(Default::default); - Some( - thread_stats - .turns_with_citations - .fetch_add(1, Ordering::Relaxed) - + 1, - ) -} diff --git a/codex-rs/ext/memories/src/contributors/citations.rs b/codex-rs/ext/memories/src/contributors/citations.rs new file mode 100644 index 0000000000..c382a5d429 --- /dev/null +++ b/codex-rs/ext/memories/src/contributors/citations.rs @@ -0,0 +1,34 @@ +use codex_extension_api::ExtensionData; +use codex_extension_api::TurnItemContributionFuture; +use codex_extension_api::TurnItemContributor; +use codex_memories_read::citations::parse_memory_citation; +use codex_protocol::items::AgentMessageContent; +use codex_protocol::items::TurnItem; +use codex_utils_stream_parser::strip_citations; + +#[derive(Debug)] +pub(crate) struct CitationContributor; + +impl TurnItemContributor for CitationContributor { + fn contribute<'a>( + &'a self, + _thread_store: &'a ExtensionData, + _turn_store: &'a ExtensionData, + item: &'a mut TurnItem, + ) -> TurnItemContributionFuture<'a> { + if let TurnItem::AgentMessage(agent_message) = item { + let text = agent_message + .content + .iter() + .map(|entry| match entry { + AgentMessageContent::Text { text } => text.as_str(), + }) + .collect::(); + let (visible_text, citations) = strip_citations(&text); + agent_message.content = vec![AgentMessageContent::Text { text: visible_text }]; + agent_message.memory_citation = parse_memory_citation(citations); + } + + Box::pin(std::future::ready(Ok(()))) + } +} diff --git a/codex-rs/ext/memories/src/contributors/mod.rs b/codex-rs/ext/memories/src/contributors/mod.rs new file mode 100644 index 0000000000..3d49a199d8 --- /dev/null +++ b/codex-rs/ext/memories/src/contributors/mod.rs @@ -0,0 +1,7 @@ +mod citations; +mod prompt; +mod tools; + +pub(crate) use citations::CitationContributor; +pub(crate) use prompt::PromptContributor; +pub(crate) use tools::ToolsContributor; diff --git a/codex-rs/ext/memories/src/contributors/prompt.rs b/codex-rs/ext/memories/src/contributors/prompt.rs new file mode 100644 index 0000000000..d15ab59963 --- /dev/null +++ b/codex-rs/ext/memories/src/contributors/prompt.rs @@ -0,0 +1,34 @@ +use codex_extension_api::ContextContributor; +use codex_extension_api::ExtensionData; +use codex_extension_api::PromptFragment; + +use crate::ctx::read_surface_enabled; + +#[derive(Debug)] +pub(crate) struct PromptContributor { + read_prompt: Option, +} + +impl PromptContributor { + pub(crate) fn new(read_prompt: Option) -> Self { + Self { read_prompt } + } +} + +impl ContextContributor for PromptContributor { + fn contribute( + &self, + _session_store: &ExtensionData, + thread_store: &ExtensionData, + ) -> Vec { + if !read_surface_enabled(thread_store) { + return Vec::new(); + } + + self.read_prompt + .as_deref() + .map(PromptFragment::developer_policy) + .into_iter() + .collect() + } +} diff --git a/codex-rs/ext/memories/src/contributors/tools.rs b/codex-rs/ext/memories/src/contributors/tools.rs new file mode 100644 index 0000000000..513a166825 --- /dev/null +++ b/codex-rs/ext/memories/src/contributors/tools.rs @@ -0,0 +1,27 @@ +use codex_extension_api::ExtensionData; +use codex_extension_api::ToolContribution; +use codex_extension_api::ToolContributor; + +use crate::ctx::read_surface_enabled; +use crate::tools::MemoriesTools; + +#[derive(Debug)] +pub(crate) struct ToolsContributor { + tools: MemoriesTools, +} + +impl ToolsContributor { + pub(crate) fn new(tools: MemoriesTools) -> Self { + Self { tools } + } +} + +impl ToolContributor for ToolsContributor { + fn tools(&self, thread_store: &ExtensionData) -> Vec { + if !read_surface_enabled(thread_store) { + return Vec::new(); + } + + self.tools.contributions() + } +} diff --git a/codex-rs/ext/memories/src/ctx.rs b/codex-rs/ext/memories/src/ctx.rs index 8c3a2fc191..c9b33370e8 100644 --- a/codex-rs/ext/memories/src/ctx.rs +++ b/codex-rs/ext/memories/src/ctx.rs @@ -1,8 +1,12 @@ -/// Runtime facts needed to decide whether read-memory surfaces are visible. -/// -/// Hosts should expose the current effective values for the thread being -/// assembled. The extension owns the policy that combines those values. -pub trait MemoriesContext { - fn memory_tool_enabled(&self) -> bool; - fn use_memories(&self) -> bool; +use codex_extension_api::ExtensionData; + +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct MemoriesReadConfig { + pub(crate) enabled: bool, +} + +pub(crate) fn read_surface_enabled(thread_store: &ExtensionData) -> bool { + thread_store + .get::() + .is_some_and(|config| config.enabled) } diff --git a/codex-rs/ext/memories/src/lib.rs b/codex-rs/ext/memories/src/lib.rs index d865ed54f8..3dcb9e49b1 100644 --- a/codex-rs/ext/memories/src/lib.rs +++ b/codex-rs/ext/memories/src/lib.rs @@ -1,36 +1,52 @@ -mod citation_output; +mod contributors; pub mod ctx; -mod list_tool; -mod prompt_contributor; -mod tool_contributor; +mod tools; use std::path::PathBuf; use std::sync::Arc; -use crate::ctx::MemoriesContext; +use crate::ctx::MemoriesReadConfig; +use codex_core::config::Config; use codex_extension_api::CodexExtension; +use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ThreadStartContributor; +use codex_features::Feature; use codex_memories_read::build_memory_tool_developer_instructions; use codex_memories_read::memory_root; -use codex_protocol::items::TurnItem; use codex_utils_absolute_path::AbsolutePathBuf; -use list_tool::ListMemoriesTool; +use contributors::CitationContributor; +use contributors::PromptContributor; +use contributors::ToolsContributor; +use tools::MemoriesTools; /// Extension that contributes memories read surfaces. #[derive(Clone, Debug)] pub struct MemoriesExtension { - read_prompt: Option, - pub(crate) list_tool: Arc, // This is just to have useful examples, it will disappear + prompt: Arc, + tools: Arc, + citations: Arc, } -impl CodexExtension for MemoriesExtension -where - C: MemoriesContext + Send + Sync + 'static, -{ - fn install(self: Arc, registry: &mut ExtensionRegistryBuilder) { - registry.tool_contributor(self.clone()); - registry.output_contributor::(self.clone()); - registry.prompt_contributor(self); +impl CodexExtension for MemoriesExtension { + fn install(self: Arc, registry: &mut ExtensionRegistryBuilder) { + registry.thread_start_contributor(self.clone()); + registry.prompt_contributor(self.prompt.clone()); + registry.tool_contributor(self.tools.clone()); + registry.turn_item_contributor(self.citations.clone()); + } +} + +impl ThreadStartContributor for MemoriesExtension { + fn contribute( + &self, + config: &Config, + _session_store: &ExtensionData, + thread_store: &ExtensionData, + ) { + thread_store.insert(MemoriesReadConfig { + enabled: config.features.enabled(Feature::MemoryTool) && config.memories.use_memories, + }); } } @@ -39,20 +55,11 @@ impl MemoriesExtension { /// it should expose. pub fn new(read_prompt: Option, memories_root: impl Into) -> Self { Self { - read_prompt, - list_tool: Arc::new(ListMemoriesTool::new(memories_root)), + prompt: Arc::new(PromptContributor::new(read_prompt)), + tools: Arc::new(ToolsContributor::new(MemoriesTools::new(memories_root))), + citations: Arc::new(CitationContributor), } } - - /// Returns the rendered developer instruction for read access, if available. - pub fn read_prompt(&self) -> Option<&str> { - self.read_prompt.as_deref() - } - - // Just for example - pub(crate) fn is_read_surface_enabled(&self, context: &C) -> bool { - context.memory_tool_enabled() && context.use_memories() - } } pub async fn extension(codex_home: &AbsolutePathBuf) -> Arc { diff --git a/codex-rs/ext/memories/src/list_tool.rs b/codex-rs/ext/memories/src/list_tool.rs deleted file mode 100644 index ae2d1d887c..0000000000 --- a/codex-rs/ext/memories/src/list_tool.rs +++ /dev/null @@ -1,323 +0,0 @@ -// 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; -use std::path::Path; -use std::path::PathBuf; -use std::pin::Pin; -use std::sync::Arc; - -use codex_extension_api::ToolCallError; -use codex_extension_api::ToolContribution; -use codex_extension_api::ToolHandler; -use codex_tools::JsonSchema; -use codex_tools::ResponsesApiTool; -use serde::Deserialize; -use serde::Serialize; -use serde_json::Value; - -const LIST_MEMORIES_TOOL_NAME: &str = "list_memories"; -const DEFAULT_LIST_MAX_RESULTS: usize = 2_000; -const MAX_LIST_RESULTS: usize = 2_000; - -#[derive(Debug)] -pub(super) struct ListMemoriesTool { - memories_root: PathBuf, -} - -impl ListMemoriesTool { - pub(super) fn new(memories_root: impl Into) -> Self { - Self { - memories_root: memories_root.into(), - } - } - - pub(super) fn contribution(self: &Arc) -> ToolContribution - where - C: Send + Sync + 'static, - { - let handler: Arc> = self.clone(); - ToolContribution::new(create_list_memories_tool(), handler).allow_parallel_calls() - } -} - -impl ToolHandler for ListMemoriesTool -where - C: Send + Sync, -{ - fn handle<'a>( - &'a self, - _context: &'a C, - arguments: Value, - ) -> Pin> + Send + 'a>> { - Box::pin(async move { - let args: ListMemoriesArgs = serde_json::from_value(arguments) - .map_err(|err| ToolCallError::new(format!("invalid arguments: {err}")))?; - tokio::fs::create_dir_all(&self.memories_root) - .await - .map_err(|err| { - ToolCallError::new(format!( - "failed to create memories root at {}: {err}", - self.memories_root.display() - )) - })?; - let response = list_memories(&self.memories_root, args) - .await - .map_err(|err| ToolCallError::new(err.to_string()))?; - serde_json::to_value(response) - .map_err(|err| ToolCallError::new(format!("failed to serialize output: {err}"))) - }) - } -} - -#[derive(Debug, Deserialize)] -#[serde(deny_unknown_fields)] -struct ListMemoriesArgs { - path: Option, - cursor: Option, - max_results: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -struct ListMemoriesResponse { - path: Option, - entries: Vec, - next_cursor: Option, - truncated: bool, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -struct MemoryEntry { - path: String, - entry_type: MemoryEntryType, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] -#[serde(rename_all = "snake_case")] -enum MemoryEntryType { - File, - Directory, -} - -#[derive(Debug, thiserror::Error)] -enum ListMemoriesError { - #[error("path '{path}' {reason}")] - InvalidPath { path: String, reason: String }, - #[error("cursor '{cursor}' {reason}")] - InvalidCursor { cursor: String, reason: String }, - #[error("path '{path}' was not found")] - NotFound { path: String }, - #[error("I/O error while reading memories: {0}")] - Io(#[from] std::io::Error), -} - -async fn list_memories( - memories_root: &Path, - args: ListMemoriesArgs, -) -> Result { - let max_results = args - .max_results - .unwrap_or(DEFAULT_LIST_MAX_RESULTS) - .min(MAX_LIST_RESULTS); - let start = resolve_scoped_path(memories_root, args.path.as_deref()).await?; - let start_index = match args.cursor.as_deref() { - Some(cursor) => cursor - .parse::() - .map_err(|_| ListMemoriesError::InvalidCursor { - cursor: cursor.to_string(), - reason: "must be a non-negative integer".to_string(), - })?, - None => 0, - }; - let Some(metadata) = metadata_or_none(&start).await? else { - return Err(ListMemoriesError::NotFound { - path: args.path.unwrap_or_default(), - }); - }; - reject_symlink(&display_relative_path(memories_root, &start), &metadata)?; - - let mut entries = if metadata.is_file() { - vec![MemoryEntry { - path: display_relative_path(memories_root, &start), - entry_type: MemoryEntryType::File, - }] - } else if metadata.is_dir() { - let mut entries = Vec::new(); - for path in read_sorted_dir_paths(&start).await? { - if is_hidden_path(&path) { - continue; - } - let Some(metadata) = metadata_or_none(&path).await? else { - continue; - }; - if metadata.file_type().is_symlink() { - continue; - } - - let entry_type = if metadata.is_dir() { - MemoryEntryType::Directory - } else if metadata.is_file() { - MemoryEntryType::File - } else { - continue; - }; - entries.push(MemoryEntry { - path: display_relative_path(memories_root, &path), - entry_type, - }); - } - entries - } else { - Vec::new() - }; - if start_index > entries.len() { - return Err(ListMemoriesError::InvalidCursor { - cursor: start_index.to_string(), - reason: "exceeds result count".to_string(), - }); - } - - let end_index = start_index.saturating_add(max_results).min(entries.len()); - let next_cursor = (end_index < entries.len()).then(|| end_index.to_string()); - let truncated = next_cursor.is_some(); - Ok(ListMemoriesResponse { - path: args.path, - entries: entries.drain(start_index..end_index).collect(), - next_cursor, - truncated, - }) -} - -async fn resolve_scoped_path( - memories_root: &Path, - relative_path: Option<&str>, -) -> Result { - let Some(relative_path) = relative_path else { - return Ok(memories_root.to_path_buf()); - }; - let relative = Path::new(relative_path); - if relative.components().any(|component| { - matches!( - component, - Component::ParentDir | Component::RootDir | Component::Prefix(_) - ) - }) { - return Err(ListMemoriesError::InvalidPath { - path: relative_path.to_string(), - reason: "must stay within the memories root".to_string(), - }); - } - if relative.components().any(is_hidden_component) { - return Err(ListMemoriesError::NotFound { - path: relative_path.to_string(), - }); - } - - let components = relative.components().collect::>(); - let mut scoped_path = memories_root.to_path_buf(); - for (index, component) in components.iter().enumerate() { - scoped_path.push(component.as_os_str()); - - let Some(metadata) = metadata_or_none(&scoped_path).await? else { - for remaining_component in components.iter().skip(index + 1) { - scoped_path.push(remaining_component.as_os_str()); - } - return Ok(scoped_path); - }; - - reject_symlink( - &display_relative_path(memories_root, &scoped_path), - &metadata, - )?; - if index + 1 < components.len() && !metadata.is_dir() { - return Err(ListMemoriesError::InvalidPath { - path: relative_path.to_string(), - reason: "traverses through a non-directory path component".to_string(), - }); - } - } - - Ok(scoped_path) -} - -async fn metadata_or_none(path: &Path) -> Result, ListMemoriesError> { - match tokio::fs::symlink_metadata(path).await { - Ok(metadata) => Ok(Some(metadata)), - Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None), - Err(err) => Err(err.into()), - } -} - -fn reject_symlink( - relative_path: &str, - metadata: &std::fs::Metadata, -) -> Result<(), ListMemoriesError> { - if metadata.file_type().is_symlink() { - return Err(ListMemoriesError::InvalidPath { - path: relative_path.to_string(), - reason: "must not be a symlink".to_string(), - }); - } - Ok(()) -} - -fn display_relative_path(root: &Path, path: &Path) -> String { - path.strip_prefix(root) - .map_or(Cow::Borrowed(path), Cow::Borrowed) - .to_string_lossy() - .replace(std::path::MAIN_SEPARATOR, "/") -} - -fn is_hidden_path(path: &Path) -> bool { - path.file_name() - .is_some_and(|name| name.to_string_lossy().starts_with('.')) -} - -fn is_hidden_component(component: Component<'_>) -> bool { - matches!(component, Component::Normal(name) if name.to_string_lossy().starts_with('.')) -} - -async fn read_sorted_dir_paths(path: &Path) -> Result, std::io::Error> { - let mut entries = tokio::fs::read_dir(path).await?; - let mut paths = Vec::new(); - while let Some(entry) = entries.next_entry().await? { - paths.push(entry.path()); - } - paths.sort(); - Ok(paths) -} - -fn create_list_memories_tool() -> ResponsesApiTool { - let properties = std::collections::BTreeMap::from([ - ( - "path".to_string(), - JsonSchema::string(Some( - "Optional relative path to list inside the Codex memories store.".to_string(), - )), - ), - ( - "cursor".to_string(), - JsonSchema::string(Some( - "Optional cursor returned by a previous list_memories call.".to_string(), - )), - ), - ( - "max_results".to_string(), - JsonSchema::integer(Some( - "Optional maximum number of entries to return.".to_string(), - )), - ), - ]); - - ResponsesApiTool { - name: LIST_MEMORIES_TOOL_NAME.to_string(), - description: - "List immediate files and directories under a path in the Codex memories store." - .to_string(), - strict: false, - defer_loading: None, - parameters: JsonSchema::object(properties, /*required*/ None, Some(false.into())), - output_schema: None, - } -} diff --git a/codex-rs/ext/memories/src/prompt_contributor.rs b/codex-rs/ext/memories/src/prompt_contributor.rs deleted file mode 100644 index 1244e2e7ce..0000000000 --- a/codex-rs/ext/memories/src/prompt_contributor.rs +++ /dev/null @@ -1,19 +0,0 @@ -use codex_extension_api::ContextContributor; -use codex_extension_api::PromptFragment; -use codex_extension_api::Stores; - -use crate::MemoriesExtension; -use crate::ctx::MemoriesContext; - -impl ContextContributor for MemoriesExtension { - fn contribute(&self, context: &C, _stores: &Stores<'_>) -> Vec { - if !self.is_read_surface_enabled(context) { - return Vec::new(); - } - - self.read_prompt() - .map(PromptFragment::developer_policy) - .into_iter() - .collect() - } -} diff --git a/codex-rs/ext/memories/src/tool_contributor.rs b/codex-rs/ext/memories/src/tool_contributor.rs deleted file mode 100644 index b4a9198d3b..0000000000 --- a/codex-rs/ext/memories/src/tool_contributor.rs +++ /dev/null @@ -1,17 +0,0 @@ -use codex_extension_api::Stores; -use codex_extension_api::ToolContribution; -use codex_extension_api::ToolContributor; - -use crate::MemoriesExtension; -use crate::ctx::MemoriesContext; - -impl ToolContributor for MemoriesExtension { - fn tools(&self, context: &C, _stores: &Stores<'_>) -> Vec> { - if !self.is_read_surface_enabled(context) { - return Vec::new(); - } - - // TODO(jif) add more tools ofc - vec![self.list_tool.contribution()] - } -} diff --git a/codex-rs/ext/memories/src/tools/list.rs b/codex-rs/ext/memories/src/tools/list.rs new file mode 100644 index 0000000000..f2d6188f44 --- /dev/null +++ b/codex-rs/ext/memories/src/tools/list.rs @@ -0,0 +1,109 @@ +use std::collections::BTreeMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use codex_extension_api::ToolCallError; +use codex_extension_api::ToolContribution; +use codex_extension_api::ToolHandler; +use codex_memories_mcp::LocalMemoriesBackend; +use codex_memories_mcp::backend::DEFAULT_LIST_MAX_RESULTS; +use codex_memories_mcp::backend::ListMemoriesRequest; +use codex_memories_mcp::backend::MAX_LIST_RESULTS; +use codex_memories_mcp::backend::MemoriesBackend; +use codex_tools::JsonSchema; +use codex_tools::ResponsesApiTool; +use serde::Deserialize; +use serde_json::Value; + +use super::backend_error; +use super::clamp_max_results; +use super::ensure_root; +use super::parse_arguments; +use super::serialize_output; + +const TOOL_NAME: &str = "list_memories"; + +#[derive(Debug)] +pub(super) struct ListMemoriesTool { + backend: LocalMemoriesBackend, +} + +impl ListMemoriesTool { + pub(super) fn new(backend: LocalMemoriesBackend) -> Self { + Self { backend } + } + + pub(super) fn contribution(self: &Arc) -> ToolContribution { + let handler: Arc = self.clone(); + ToolContribution::new(tool_spec(), handler).allow_parallel_calls() + } +} + +impl ToolHandler for ListMemoriesTool { + fn handle<'a>( + &'a self, + arguments: Value, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + let args: ListArgs = parse_arguments(arguments)?; + ensure_root(&self.backend).await?; + let response = self + .backend + .list(ListMemoriesRequest { + path: args.path, + cursor: args.cursor, + max_results: clamp_max_results( + args.max_results, + DEFAULT_LIST_MAX_RESULTS, + MAX_LIST_RESULTS, + ), + }) + .await + .map_err(backend_error)?; + serialize_output(response) + }) + } +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct ListArgs { + path: Option, + cursor: Option, + max_results: Option, +} + +fn tool_spec() -> ResponsesApiTool { + let properties = BTreeMap::from([ + ( + "path".to_string(), + JsonSchema::string(Some( + "Optional relative path to list inside the Codex memories store.".to_string(), + )), + ), + ( + "cursor".to_string(), + JsonSchema::string(Some( + "Optional cursor returned by a previous list_memories call.".to_string(), + )), + ), + ( + "max_results".to_string(), + JsonSchema::integer(Some( + "Optional maximum number of entries to return.".to_string(), + )), + ), + ]); + + ResponsesApiTool { + name: TOOL_NAME.to_string(), + description: + "List immediate files and directories under a path in the Codex memories store." + .to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object(properties, /*required*/ None, Some(false.into())), + output_schema: None, + } +} diff --git a/codex-rs/ext/memories/src/tools/mod.rs b/codex-rs/ext/memories/src/tools/mod.rs new file mode 100644 index 0000000000..459c00f589 --- /dev/null +++ b/codex-rs/ext/memories/src/tools/mod.rs @@ -0,0 +1,73 @@ +mod list; +mod read; +mod search; + +use std::fmt::Display; +use std::path::PathBuf; +use std::sync::Arc; + +use codex_extension_api::ToolCallError; +use codex_extension_api::ToolContribution; +use codex_memories_mcp::LocalMemoriesBackend; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; + +use self::list::ListMemoriesTool; +use self::read::ReadMemoryTool; +use self::search::SearchMemoriesTool; + +#[derive(Debug)] +pub(crate) struct MemoriesTools { + list: Arc, + read: Arc, + search: Arc, +} + +impl MemoriesTools { + pub(crate) fn new(memories_root: impl Into) -> Self { + let backend = LocalMemoriesBackend::from_memory_root(memories_root); + Self { + list: Arc::new(ListMemoriesTool::new(backend.clone())), + read: Arc::new(ReadMemoryTool::new(backend.clone())), + search: Arc::new(SearchMemoriesTool::new(backend)), + } + } + + pub(crate) fn contributions(&self) -> Vec { + vec![ + self.list.contribution(), + self.read.contribution(), + self.search.contribution(), + ] + } +} + +async fn ensure_root(backend: &LocalMemoriesBackend) -> Result<(), ToolCallError> { + tokio::fs::create_dir_all(backend.root()) + .await + .map_err(|err| { + ToolCallError::new(format!( + "failed to create memories root at {}: {err}", + backend.root().display() + )) + }) +} + +fn parse_arguments(arguments: Value) -> Result { + serde_json::from_value(arguments) + .map_err(|err| ToolCallError::new(format!("invalid arguments: {err}"))) +} + +fn serialize_output(output: T) -> Result { + serde_json::to_value(output) + .map_err(|err| ToolCallError::new(format!("failed to serialize output: {err}"))) +} + +fn backend_error(err: impl Display) -> ToolCallError { + ToolCallError::new(err.to_string()) +} + +fn clamp_max_results(requested: Option, default: usize, max: usize) -> usize { + requested.unwrap_or(default).clamp(1, max) +} diff --git a/codex-rs/ext/memories/src/tools/read.rs b/codex-rs/ext/memories/src/tools/read.rs new file mode 100644 index 0000000000..cbdd74b08e --- /dev/null +++ b/codex-rs/ext/memories/src/tools/read.rs @@ -0,0 +1,108 @@ +use std::collections::BTreeMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use codex_extension_api::ToolCallError; +use codex_extension_api::ToolContribution; +use codex_extension_api::ToolHandler; +use codex_memories_mcp::LocalMemoriesBackend; +use codex_memories_mcp::backend::DEFAULT_READ_MAX_TOKENS; +use codex_memories_mcp::backend::MemoriesBackend; +use codex_memories_mcp::backend::ReadMemoryRequest; +use codex_tools::JsonSchema; +use codex_tools::ResponsesApiTool; +use serde::Deserialize; +use serde_json::Value; + +use super::backend_error; +use super::ensure_root; +use super::parse_arguments; +use super::serialize_output; + +const TOOL_NAME: &str = "read_memory"; + +#[derive(Debug)] +pub(super) struct ReadMemoryTool { + backend: LocalMemoriesBackend, +} + +impl ReadMemoryTool { + pub(super) fn new(backend: LocalMemoriesBackend) -> Self { + Self { backend } + } + + pub(super) fn contribution(self: &Arc) -> ToolContribution { + let handler: Arc = self.clone(); + ToolContribution::new(tool_spec(), handler).allow_parallel_calls() + } +} + +impl ToolHandler for ReadMemoryTool { + fn handle<'a>( + &'a self, + arguments: Value, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + let args: ReadArgs = parse_arguments(arguments)?; + ensure_root(&self.backend).await?; + let response = self + .backend + .read(ReadMemoryRequest { + path: args.path, + line_offset: args.line_offset.unwrap_or(1), + max_lines: args.max_lines, + max_tokens: DEFAULT_READ_MAX_TOKENS, + }) + .await + .map_err(backend_error)?; + serialize_output(response) + }) + } +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct ReadArgs { + path: String, + line_offset: Option, + max_lines: Option, +} + +fn tool_spec() -> ResponsesApiTool { + let properties = BTreeMap::from([ + ( + "path".to_string(), + JsonSchema::string(Some( + "Relative path of the Codex memory file to read.".to_string(), + )), + ), + ( + "line_offset".to_string(), + JsonSchema::integer(Some( + "Optional 1-indexed line number to start reading from.".to_string(), + )), + ), + ( + "max_lines".to_string(), + JsonSchema::integer(Some( + "Optional maximum number of lines to return.".to_string(), + )), + ), + ]); + + ResponsesApiTool { + name: TOOL_NAME.to_string(), + description: + "Read a Codex memory file by relative path, optionally starting at a 1-indexed line offset and limiting the number of lines returned." + .to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object( + properties, + Some(vec!["path".to_string()]), + Some(false.into()), + ), + output_schema: None, + } +} diff --git a/codex-rs/ext/memories/src/tools/search.rs b/codex-rs/ext/memories/src/tools/search.rs new file mode 100644 index 0000000000..3ba154901a --- /dev/null +++ b/codex-rs/ext/memories/src/tools/search.rs @@ -0,0 +1,192 @@ +use std::collections::BTreeMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use codex_extension_api::ToolCallError; +use codex_extension_api::ToolContribution; +use codex_extension_api::ToolHandler; +use codex_memories_mcp::LocalMemoriesBackend; +use codex_memories_mcp::backend::DEFAULT_SEARCH_MAX_RESULTS; +use codex_memories_mcp::backend::MAX_SEARCH_RESULTS; +use codex_memories_mcp::backend::MemoriesBackend; +use codex_memories_mcp::backend::SearchMatchMode; +use codex_memories_mcp::backend::SearchMemoriesRequest; +use codex_tools::JsonSchema; +use codex_tools::ResponsesApiTool; +use serde::Deserialize; +use serde_json::Value; +use serde_json::json; + +use super::backend_error; +use super::clamp_max_results; +use super::ensure_root; +use super::parse_arguments; +use super::serialize_output; + +const TOOL_NAME: &str = "search_memories"; + +#[derive(Debug)] +pub(super) struct SearchMemoriesTool { + backend: LocalMemoriesBackend, +} + +impl SearchMemoriesTool { + pub(super) fn new(backend: LocalMemoriesBackend) -> Self { + Self { backend } + } + + pub(super) fn contribution(self: &Arc) -> ToolContribution { + let handler: Arc = self.clone(); + ToolContribution::new(tool_spec(), handler).allow_parallel_calls() + } +} + +impl ToolHandler for SearchMemoriesTool { + fn handle<'a>( + &'a self, + arguments: Value, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + let args: SearchArgs = parse_arguments(arguments)?; + ensure_root(&self.backend).await?; + let response = self + .backend + .search(args.into_request()) + .await + .map_err(backend_error)?; + serialize_output(response) + }) + } +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct SearchArgs { + queries: Vec, + match_mode: Option, + path: Option, + cursor: Option, + context_lines: Option, + case_sensitive: Option, + normalized: Option, + max_results: Option, +} + +impl SearchArgs { + fn into_request(self) -> SearchMemoriesRequest { + SearchMemoriesRequest { + queries: self.queries, + match_mode: self.match_mode.unwrap_or(SearchMatchMode::Any), + path: self.path, + cursor: self.cursor, + context_lines: self.context_lines.unwrap_or(0), + case_sensitive: self.case_sensitive.unwrap_or(true), + normalized: self.normalized.unwrap_or(false), + max_results: clamp_max_results( + self.max_results, + DEFAULT_SEARCH_MAX_RESULTS, + MAX_SEARCH_RESULTS, + ), + } + } +} + +fn tool_spec() -> ResponsesApiTool { + let properties = BTreeMap::from([ + ( + "queries".to_string(), + JsonSchema::array( + JsonSchema::string(None), + Some("Substrings to search for.".to_string()), + ), + ), + ("match_mode".to_string(), search_match_mode_schema()), + ( + "path".to_string(), + JsonSchema::string(Some( + "Optional relative path to search inside the Codex memories store.".to_string(), + )), + ), + ( + "cursor".to_string(), + JsonSchema::string(Some( + "Optional cursor returned by a previous search_memories call.".to_string(), + )), + ), + ( + "context_lines".to_string(), + JsonSchema::integer(Some( + "Optional number of context lines to include around each match.".to_string(), + )), + ), + ( + "case_sensitive".to_string(), + JsonSchema::boolean(Some( + "Whether matching should be case-sensitive.".to_string(), + )), + ), + ( + "normalized".to_string(), + JsonSchema::boolean(Some( + "Whether to normalize separators while matching.".to_string(), + )), + ), + ( + "max_results".to_string(), + JsonSchema::integer(Some( + "Optional maximum number of matches to return.".to_string(), + )), + ), + ]); + + ResponsesApiTool { + name: TOOL_NAME.to_string(), + description: + "Search Codex memory files for substring matches, optionally normalizing separators or requiring all query substrings on the same line or within a line window." + .to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object( + properties, + Some(vec!["queries".to_string()]), + Some(false.into()), + ), + output_schema: None, + } +} + +fn search_match_mode_schema() -> JsonSchema { + JsonSchema::any_of( + vec![ + tagged_match_mode_schema("any", None), + tagged_match_mode_schema("all_on_same_line", None), + tagged_match_mode_schema( + "all_within_lines", + Some(( + "line_count".to_string(), + JsonSchema::integer(Some( + "Positive number of lines in the matching window.".to_string(), + )), + )), + ), + ], + Some("How multiple queries must match.".to_string()), + ) +} + +fn tagged_match_mode_schema( + name: &str, + extra_property: Option<(String, JsonSchema)>, +) -> JsonSchema { + let mut properties = BTreeMap::from([( + "type".to_string(), + JsonSchema::string_enum(vec![json!(name)], None), + )]); + let mut required = vec!["type".to_string()]; + if let Some((name, schema)) = extra_property { + required.push(name.clone()); + properties.insert(name, schema); + } + JsonSchema::object(properties, Some(required), Some(false.into())) +} diff --git a/codex-rs/mcp-server/Cargo.toml b/codex-rs/mcp-server/Cargo.toml index 74873023da..e2692a6639 100644 --- a/codex-rs/mcp-server/Cargo.toml +++ b/codex-rs/mcp-server/Cargo.toml @@ -21,6 +21,7 @@ codex-arg0 = { workspace = true } codex-config = { workspace = true } codex-core = { workspace = true } codex-exec-server = { workspace = true } +codex-extension-api = { workspace = true } codex-login = { workspace = true } codex-protocol = { workspace = true } codex-utils-cli = { workspace = true } diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index d64fc43b1b..bb6d4bb84a 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -7,6 +7,7 @@ use codex_core::ThreadManager; use codex_core::config::Config; use codex_core::thread_store_from_config; use codex_exec_server::EnvironmentManager; +use codex_extension_api::empty_extension_registry; use codex_login::AuthManager; use codex_login::default_client::USER_AGENT_SUFFIX; use codex_login::default_client::get_codex_user_agent; @@ -68,6 +69,7 @@ impl MessageProcessor { auth_manager, SessionSource::Mcp, environment_manager, + empty_extension_registry(), /*analytics_events_client*/ None, thread_store_from_config(config.as_ref(), state_db.clone()), state_db.clone(), diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 6817f677e6..335d30595d 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -54,6 +54,7 @@ use codex_core_api::UserInput; use codex_core_api::WebSearchMode; use codex_core_api::arg0_dispatch_or_else; use codex_core_api::built_in_model_providers; +use codex_core_api::empty_extension_registry; use codex_core_api::find_codex_home; use codex_core_api::init_state_db; use codex_core_api::item_event_to_server_notification; @@ -122,6 +123,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { auth_manager, SessionSource::Exec, environment_manager, + empty_extension_registry(), /*analytics_events_client*/ None, Arc::clone(&thread_store), state_db,