From cffcb59f5b5322c523f122b3f036b6da9507731b Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 7 May 2026 15:39:46 +0100 Subject: [PATCH] v4 --- codex-rs/Cargo.lock | 16 +++++++- codex-rs/Cargo.toml | 2 + .../examples/enabled_extensions.rs | 2 +- .../ext/extension-api/src/contributors.rs | 25 +++++++----- codex-rs/ext/extension-api/src/lib.rs | 2 + codex-rs/ext/extension-api/src/registry.rs | 40 +++++++++++++++++++ codex-rs/ext/memories/Cargo.toml | 2 + codex-rs/ext/memories/src/citation_output.rs | 31 ++++++++++++++ codex-rs/ext/memories/src/lib.rs | 3 ++ 9 files changed, 110 insertions(+), 13 deletions(-) create mode 100644 codex-rs/ext/memories/src/citation_output.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 42b1a89b68..19ef21197d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2444,12 +2444,14 @@ dependencies = [ "codex-core-skills", "codex-exec-server", "codex-execpolicy", + "codex-extension-api", "codex-features", "codex-feedback", "codex-git-utils", "codex-hooks", "codex-login", "codex-mcp", + "codex-memories", "codex-memories-read", "codex-model-provider", "codex-model-provider-info", @@ -2785,9 +2787,12 @@ name = "codex-extension-api" version = "0.0.0" dependencies = [ "codex-git-attribution", + "codex-guardian", "codex-memories", "codex-multi-agent-v2", - "rmcp", + "codex-tools", + "serde_json", + "thiserror 2.0.18", ] [[package]] @@ -2899,6 +2904,13 @@ dependencies = [ "walkdir", ] +[[package]] +name = "codex-guardian" +version = "0.0.0" +dependencies = [ + "codex-extension-api", +] + [[package]] name = "codex-hooks" version = "0.0.0" @@ -3091,6 +3103,8 @@ dependencies = [ "codex-utils-absolute-path", "rmcp", "serde_json", + "thiserror 2.0.18", + "tokio", ] [[package]] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index f461f16eb0..5c2572f16f 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -46,6 +46,7 @@ members = [ "execpolicy", "execpolicy-legacy", "ext/extension-api", + "ext/guardian", "ext/git-attribution", "ext/memories", "ext/multi-agent-v2", @@ -168,6 +169,7 @@ 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" } diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions.rs b/codex-rs/ext/extension-api/examples/enabled_extensions.rs index cc12bd9f8b..ab92f4acb5 100644 --- a/codex-rs/ext/extension-api/examples/enabled_extensions.rs +++ b/codex-rs/ext/extension-api/examples/enabled_extensions.rs @@ -71,7 +71,7 @@ fn main() { mod ctx { use codex_git_attribution::GitAttributionContext; use codex_guardian::GuardianContext; - use codex_memories::MemoriesContext; + use codex_memories::ctx::MemoriesContext; use codex_multi_agent_v2::MultiAgentV2Context; use codex_multi_agent_v2::UsageHintAudience; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 3e57cea152..da22b0611f 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -5,9 +5,11 @@ //! vocabulary from one place, while each contribution family keeps its own //! supporting types nearby. +use std::future::Future; + mod prompt; -mod tool; mod session_lifecycle; +mod tool; pub use prompt::PromptFragment; pub use prompt::PromptSlot; @@ -27,18 +29,19 @@ pub trait ToolContributor: Send + Sync { fn tools(&self, context: &C) -> Vec>; } -/// Analyze or perform computation on the output. Works only if the post-processing pipeline is ordered -pub type OutputContributionFuture<'a> = std::pin::Pin> + Send + 'a>>; +/// Future returned by one ordered output contribution. +pub type OutputContributionFuture<'a> = + std::pin::Pin> + Send + 'a>>; + +/// Ordered post-processing contribution for one completed output value. +/// +/// Implementations may inspect or mutate `output`; hosts are expected to run +/// contributors sequentially so each contributor observes the result of the +/// previous one. pub trait OutputContributor: Send + Sync { - fn contribute<'a>( - &'a self, - context: &'a C, - output: &'a mut O, - ) -> OutputContributionFuture<'a>; + fn contribute<'a>(&'a self, context: &'a C, output: &'a mut O) -> OutputContributionFuture<'a>; } - - // TODO: WIP /// 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?) @@ -52,4 +55,4 @@ pub trait SessionLifecycleContributor: Send + Sync { &self, event: session_lifecycle::Event<'_, S, T>, ) -> impl Future + Send; -} \ No newline at end of file +} diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 571d706f76..279e5e1f25 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -13,6 +13,8 @@ mod registry; pub use contributors::ApprovalInterceptorContributor; pub use contributors::ContextContributor; +pub use contributors::OutputContributionFuture; +pub use contributors::OutputContributor; pub use contributors::PromptFragment; pub use contributors::PromptSlot; pub use contributors::ToolCallError; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index ca70e5e110..8de1473715 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -1,14 +1,19 @@ +use std::any::Any; +use std::any::TypeId; +use std::collections::HashMap; use std::sync::Arc; use crate::ApprovalInterceptorContributor; use crate::CodexExtension; use crate::ContextContributor; +use crate::OutputContributor; use crate::ToolContributor; /// Mutable registry used while extensions install their typed contributions. pub struct ExtensionRegistryBuilder { context_contributors: Vec>>, tool_contributors: Vec>>, + output_contributors: HashMap>, approval_interceptor_contributors: Vec>>, } @@ -18,6 +23,7 @@ impl Default for ExtensionRegistryBuilder { approval_interceptor_contributors: Vec::new(), context_contributors: Vec::new(), tool_contributors: Vec::new(), + output_contributors: HashMap::new(), } } } @@ -64,12 +70,30 @@ impl ExtensionRegistryBuilder { self.tool_contributors.push(contributor); } + /// Registers one ordered output contributor for output type `O`. + pub fn output_contributor(&mut self, contributor: Arc>) + where + C: 'static, + O: 'static, + { + let Some(contributors) = self + .output_contributors + .entry(TypeId::of::()) + .or_insert_with(|| Box::new(Vec::>>::new())) + .downcast_mut::>>>() + else { + unreachable!("output contributor bucket type must match its registered output type"); + }; + contributors.push(contributor); + } + /// Finishes construction and returns the immutable registry. pub fn build(self) -> ExtensionRegistry { ExtensionRegistry { approval_interceptor_contributors: self.approval_interceptor_contributors, prompt_contributors: self.context_contributors, tool_contributors: self.tool_contributors, + output_contributors: self.output_contributors, } } } @@ -79,6 +103,7 @@ pub struct ExtensionRegistry { approval_interceptor_contributors: Vec>>, prompt_contributors: Vec>>, tool_contributors: Vec>>, + output_contributors: HashMap>, } impl ExtensionRegistry { @@ -98,4 +123,19 @@ impl ExtensionRegistry { pub fn tool_contributors(&self) -> &[Arc>] { &self.tool_contributors } + + /// Returns the registered ordered output contributors for output type `O`. + pub fn output_contributors(&self) -> &[Arc>] + where + C: 'static, + O: 'static, + { + self.output_contributors + .get(&TypeId::of::()) + .and_then(|contributors| { + contributors.downcast_ref::>>>() + }) + .map(Vec::as_slice) + .unwrap_or_default() + } } diff --git a/codex-rs/ext/memories/Cargo.toml b/codex-rs/ext/memories/Cargo.toml index 9275dfbe79..cb133393e3 100644 --- a/codex-rs/ext/memories/Cargo.toml +++ b/codex-rs/ext/memories/Cargo.toml @@ -14,8 +14,10 @@ workspace = true [dependencies] codex-extension-api = { workspace = true } codex-memories-read = { workspace = true } +codex-protocol = { workspace = true } codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } +codex-utils-stream-parser = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } thiserror = { workspace = true } diff --git a/codex-rs/ext/memories/src/citation_output.rs b/codex-rs/ext/memories/src/citation_output.rs new file mode 100644 index 0000000000..4d00a1df26 --- /dev/null +++ b/codex-rs/ext/memories/src/citation_output.rs @@ -0,0 +1,31 @@ +use codex_extension_api::OutputContributionFuture; +use codex_extension_api::OutputContributor; +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; + +impl OutputContributor for MemoriesExtension { + fn contribute<'a>( + &'a self, + _context: &'a C, + 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); + } + + Box::pin(std::future::ready(Ok(()))) + } +} diff --git a/codex-rs/ext/memories/src/lib.rs b/codex-rs/ext/memories/src/lib.rs index fcc33bed5d..6f8334ea00 100644 --- a/codex-rs/ext/memories/src/lib.rs +++ b/codex-rs/ext/memories/src/lib.rs @@ -2,6 +2,7 @@ #![forbid(unsafe_code)] +mod citation_output; pub mod ctx; mod list_tool; @@ -17,6 +18,7 @@ use codex_extension_api::ToolContribution; use codex_extension_api::ToolContributor; use codex_memories_read::build_memory_tool_developer_instructions; use codex_memories_read::memory_root; +use codex_protocol::items::TurnItem; use codex_utils_absolute_path::AbsolutePathBuf; use list_tool::ListMemoriesTool; @@ -53,6 +55,7 @@ impl ContextContributor for MemoriesExtension { impl CodexExtension for MemoriesExtension { fn install(self: Arc, registry: &mut ExtensionRegistryBuilder) { registry.tool_contributor(self.clone()); + registry.output_contributor::(self.clone()); registry.prompt_contributor(self); } }