diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ec654051e4..76b6209813 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2816,6 +2816,16 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "codex-extension-api" +version = "0.0.0" +dependencies = [ + "codex-protocol", + "codex-tools", + "serde_json", + "thiserror 2.0.18", +] + [[package]] name = "codex-external-agent-migration" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index dd0cd7e491..81a8528a14 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -45,6 +45,7 @@ members = [ "exec-server", "execpolicy", "execpolicy-legacy", + "ext/extension-api", "external-agent-migration", "external-agent-sessions", "keyring-store", @@ -160,6 +161,7 @@ codex-exec = { path = "exec" } codex-file-system = { path = "file-system" } codex-exec-server = { path = "exec-server" } codex-execpolicy = { path = "execpolicy" } +codex-extension-api = { path = "ext/extension-api" } 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" } @@ -468,6 +470,7 @@ unwrap_used = "deny" [workspace.metadata.cargo-shear] ignored = [ "codex-agent-graph-store", + "codex-extension-api", "icu_provider", "openssl-sys", "codex-v8-poc", diff --git a/codex-rs/ext/extension-api/BUILD.bazel b/codex-rs/ext/extension-api/BUILD.bazel new file mode 100644 index 0000000000..c79ad601a4 --- /dev/null +++ b/codex-rs/ext/extension-api/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "extension-api", + crate_name = "codex_extension_api", +) diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml new file mode 100644 index 0000000000..9ff3595f03 --- /dev/null +++ b/codex-rs/ext/extension-api/Cargo.toml @@ -0,0 +1,20 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-extension-api" +version.workspace = true + +[lib] +name = "codex_extension_api" +path = "src/lib.rs" +test = false +doctest = false + +[lints] +workspace = true + +[dependencies] +codex-protocol = { workspace = true } +codex-tools = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } diff --git a/codex-rs/ext/extension-api/examples/enabled_extensions.rs b/codex-rs/ext/extension-api/examples/enabled_extensions.rs new file mode 100644 index 0000000000..f091aa7ab5 --- /dev/null +++ b/codex-rs/ext/extension-api/examples/enabled_extensions.rs @@ -0,0 +1,68 @@ +#[path = "enabled_extensions/shared_state_extension.rs"] +mod shared_state_extension; + +use std::sync::Arc; + +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionRegistryBuilder; +use shared_state_extension::SharedStateExtension; +use shared_state_extension::recorded_style_contributions; +use shared_state_extension::recorded_usage_contributions; + +fn main() { + // 1. Build the extension value owned by the host. + let extension = Arc::new(SharedStateExtension); + + // 2. Install it into the registry for the thread-start input type this host exposes. + let registry = ExtensionRegistryBuilder::<()>::new() + .with_extension(extension) + .build(); + + // 3. The host decides which stores are shared. + let session_store = ExtensionData::new(); + let first_thread_store = ExtensionData::new(); + let second_thread_store = ExtensionData::new(); + + // 4. Reusing the same session store shares session state across threads. + let first_thread_fragments = contribute_prompt(®istry, &session_store, &first_thread_store); + contribute_prompt(®istry, &session_store, &first_thread_store); + contribute_prompt(®istry, &session_store, &second_thread_store); + + println!("first prompt fragments: {}", first_thread_fragments.len()); + println!( + "session style contributions: {}", + recorded_style_contributions(&session_store) + ); + println!( + "session usage contributions: {}", + recorded_usage_contributions(&session_store) + ); + println!( + "first thread style contributions: {}", + recorded_style_contributions(&first_thread_store) + ); + println!( + "first thread usage contributions: {}", + recorded_usage_contributions(&first_thread_store) + ); + println!( + "second thread style contributions: {}", + recorded_style_contributions(&second_thread_store) + ); + println!( + "second thread usage contributions: {}", + recorded_usage_contributions(&second_thread_store) + ); +} + +fn contribute_prompt( + registry: &codex_extension_api::ExtensionRegistry<()>, + session_store: &ExtensionData, + thread_store: &ExtensionData, +) -> Vec { + registry + .context_contributors() + .iter() + .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 new file mode 100644 index 0000000000..9b6612e3ae --- /dev/null +++ b/codex-rs/ext/extension-api/examples/enabled_extensions/shared_state_extension.rs @@ -0,0 +1,100 @@ +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_extension_api::CodexExtension; +use codex_extension_api::ContextContributor; +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::PromptFragment; + +/// Small tutorial extension that installs two prompt contributors. +#[derive(Debug, Default)] +pub struct SharedStateExtension; + +impl CodexExtension<()> for SharedStateExtension { + fn install(self: Arc, registry: &mut ExtensionRegistryBuilder<()>) { + registry.prompt_contributor(Arc::new(StyleContributor)); + registry.prompt_contributor(Arc::new(UsageContributor)); + } +} + +#[derive(Debug)] +struct StyleContributor; + +impl ContextContributor for StyleContributor { + fn contribute( + &self, + session_store: &ExtensionData, + thread_store: &ExtensionData, + ) -> Vec { + contribution_counts(session_store).record_style(); + contribution_counts(thread_store).record_style(); + + vec![PromptFragment::developer_policy( + "Prefer short answers unless the user asks for detail.", + )] + } +} + +#[derive(Debug)] +struct UsageContributor; + +impl ContextContributor for UsageContributor { + fn contribute( + &self, + session_store: &ExtensionData, + thread_store: &ExtensionData, + ) -> Vec { + contribution_counts(session_store).record_usage(); + contribution_counts(thread_store).record_usage(); + + vec![PromptFragment::developer_capability( + "This extension can contribute more than one prompt fragment.", + )] + } +} + +/// Returns how many style contributions were recorded in `store`. +pub fn recorded_style_contributions(store: &ExtensionData) -> u64 { + store + .get::() + .map(|counts| counts.style()) + .unwrap_or_default() +} + +/// Returns how many usage contributions were recorded in `store`. +pub fn recorded_usage_contributions(store: &ExtensionData) -> u64 { + store + .get::() + .map(|counts| counts.usage()) + .unwrap_or_default() +} + +#[derive(Debug, Default)] +struct ContributionCounts { + style: AtomicU64, + usage: AtomicU64, +} + +impl ContributionCounts { + fn record_style(&self) { + self.style.fetch_add(1, Ordering::Relaxed); + } + + fn record_usage(&self) { + self.usage.fetch_add(1, Ordering::Relaxed); + } + + fn style(&self) -> u64 { + self.style.load(Ordering::Relaxed) + } + + fn usage(&self) -> u64 { + self.usage.load(Ordering::Relaxed) + } +} + +fn contribution_counts(store: &ExtensionData) -> Arc { + store.get_or_init::(Default::default) +} diff --git a/codex-rs/ext/extension-api/notes.md b/codex-rs/ext/extension-api/notes.md new file mode 100644 index 0000000000..e73b106f6a --- /dev/null +++ b/codex-rs/ext/extension-api/notes.md @@ -0,0 +1,14 @@ +Everything becomes a good contributor design, which contributors do we need? + +git attribution Context +memories Context + Tool + Output +guardian Context + Request +goal Tool + Runtime +image generation Tool + Output +skills Context + Turn +personality Context +plugins / apps / connectors Context + Turn +shell snapshot Runtime +web search Tool +AGENTS.md Context (Runtime too only if you want eager refresh/cache behavior) +future sandboxing probably Request + Runtime diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs new file mode 100644 index 0000000000..34096c9baf --- /dev/null +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -0,0 +1,65 @@ +use std::future::Future; + +use codex_protocol::items::TurnItem; + +use crate::ExtensionData; + +mod prompt; +mod tool; + +pub use prompt::PromptFragment; +pub use prompt::PromptSlot; +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 { + fn contribute( + &self, + session_store: &ExtensionData, + thread_store: &ExtensionData, + ) -> Vec; +} + +/// Extension contribution that exposes native tools owned by a feature. +pub trait ToolContributor: Send + Sync { + /// Returns the native tools visible for the supplied runtime context. + fn tools(&self, thread_store: &ExtensionData) -> Vec; +} + +/// Future returned by one ordered turn-item contribution. +pub type TurnItemContributionFuture<'a> = + std::pin::Pin> + Send + 'a>>; + +/// Ordered post-processing contribution for one parsed turn item. +/// +/// 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 { + fn contribute<'a>( + &'a self, + thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + item: &'a mut TurnItem, + ) -> TurnItemContributionFuture<'a>; +} + +// 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 { + /// Returns whether this contributor should intercept approvals in `context`. + fn intercepts_approvals( + &self, + thread_store: &ExtensionData, + turn_store: &ExtensionData, + ) -> bool; +} diff --git a/codex-rs/ext/extension-api/src/contributors/prompt.rs b/codex-rs/ext/extension-api/src/contributors/prompt.rs new file mode 100644 index 0000000000..1e29b72f9e --- /dev/null +++ b/codex-rs/ext/extension-api/src/contributors/prompt.rs @@ -0,0 +1,50 @@ +// All this file should be replaced by the existing fragment implementation ofc + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum PromptSlot { + DeveloperPolicy, + DeveloperCapabilities, + ContextualUser, + SeparateDeveloper, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PromptFragment { + slot: PromptSlot, + text: String, +} + +impl PromptFragment { + /// Creates a prompt fragment for the given slot. + pub fn new(slot: PromptSlot, text: impl Into) -> Self { + Self { + slot, + text: text.into(), + } + } + + /// Creates a developer-policy prompt fragment. + pub fn developer_policy(text: impl Into) -> Self { + Self::new(PromptSlot::DeveloperPolicy, text) + } + + /// Creates a developer-capabilities prompt fragment. + pub fn developer_capability(text: impl Into) -> Self { + Self::new(PromptSlot::DeveloperCapabilities, text) + } + + /// Creates a separate top-level developer prompt fragment. + pub fn separate_developer(text: impl Into) -> Self { + Self::new(PromptSlot::SeparateDeveloper, text) + } + + /// Returns the target prompt slot. + pub fn slot(&self) -> PromptSlot { + self.slot + } + + /// Returns the model-visible text. + pub fn text(&self) -> &str { + &self.text + } +} diff --git a/codex-rs/ext/extension-api/src/contributors/tool.rs b/codex-rs/ext/extension-api/src/contributors/tool.rs new file mode 100644 index 0000000000..b40e29a8e2 --- /dev/null +++ b/codex-rs/ext/extension-api/src/contributors/tool.rs @@ -0,0 +1,68 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use codex_tools::ResponsesApiTool; +use serde_json::Value; +use thiserror::Error; + +// TMP while we don't have the fully extracted tools +#[derive(Clone)] +pub struct ToolContribution { + spec: ResponsesApiTool, + handler: Arc, + supports_parallel_tool_calls: bool, +} + +impl ToolContribution { + pub fn new(spec: ResponsesApiTool, handler: Arc) -> Self { + Self { + spec, + handler, + supports_parallel_tool_calls: false, + } + } + + #[must_use] + pub fn allow_parallel_calls(mut self) -> Self { + self.supports_parallel_tool_calls = true; + self + } + + pub fn spec(&self) -> &ResponsesApiTool { + &self.spec + } + + pub fn supports_parallel_tool_calls(&self) -> bool { + self.supports_parallel_tool_calls + } + + pub fn handler(&self) -> Arc { + Arc::clone(&self.handler) + } +} + +//////// Just to make it compile //////////////////////////////// +pub trait ToolHandler: Send + Sync { + /// Handles one JSON-encoded invocation for this tool. + fn handle<'a>( + &'a self, + arguments: Value, + ) -> Pin> + Send + 'a>>; +} + +/// Error returned by a contributed native tool handler. +#[derive(Clone, Debug, Error, PartialEq, Eq)] +#[error("{message}")] +pub struct ToolCallError { + message: String, +} + +impl ToolCallError { + /// Creates a contributed-tool error with the supplied model-visible text. + pub fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } +} diff --git a/codex-rs/ext/extension-api/src/extension.rs b/codex-rs/ext/extension-api/src/extension.rs new file mode 100644 index 0000000000..c38f22e9cb --- /dev/null +++ b/codex-rs/ext/extension-api/src/extension.rs @@ -0,0 +1,12 @@ +use std::sync::Arc; + +use crate::ExtensionRegistryBuilder; + +/// First-party extension that can install one or more typed runtime contributions. +/// +/// Implementations should use [`Self::install`] only to register the concrete +/// providers they own. +pub trait CodexExtension: Send + Sync { + /// Registers this extension's concrete typed contributions. + fn install(self: Arc, registry: &mut ExtensionRegistryBuilder); +} diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs new file mode 100644 index 0000000000..0b07f9e285 --- /dev/null +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -0,0 +1,21 @@ +mod contributors; +mod extension; +mod registry; +mod state; + +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; +pub use contributors::ToolHandler; +pub use contributors::TurnItemContributionFuture; +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 new file mode 100644 index 0000000000..4a90e2e606 --- /dev/null +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -0,0 +1,134 @@ +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 { + 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(), + turn_item_contributors: Vec::new(), + } + } +} + +impl ExtensionRegistryBuilder { + /// Creates an empty registry builder. + pub fn new() -> Self { + Self::default() + } + + /// Installs one extension and returns the builder. + #[must_use] + pub fn with_extension(mut self, extension: Arc) -> Self + where + E: CodexExtension + 'static, + { + self.install_extension(extension); + self + } + + /// Installs one extension into the registry under construction. + pub fn install_extension(&mut self, extension: Arc) + where + E: CodexExtension + 'static, + { + extension.install(self); + } + + /// Registers one approval interceptor contributor. + pub fn approval_interceptor_contributor( + &mut self, + 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) { + self.context_contributors.push(contributor); + } + + /// Registers one native tool contributor. + 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) { + 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, + turn_item_contributors: self.turn_item_contributors, + } + } +} + +/// Immutable typed registry produced after extensions are installed. +pub struct ExtensionRegistry { + 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] { + &self.approval_interceptor_contributors + } + + /// Returns the registered prompt contributors. + pub fn context_contributors(&self) -> &[Arc] { + &self.context_contributors + } + + /// Returns the registered native tool contributors. + pub fn tool_contributors(&self) -> &[Arc] { + &self.tool_contributors + } + + /// Returns the registered ordered turn-item contributors. + 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/extension-api/src/state.rs b/codex-rs/ext/extension-api/src/state.rs new file mode 100644 index 0000000000..2fc9a8725f --- /dev/null +++ b/codex-rs/ext/extension-api/src/state.rs @@ -0,0 +1,77 @@ +use std::any::Any; +use std::any::TypeId; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; + +type ErasedData = Arc; + +/// Typed extension-owned data attached to one host object. +#[derive(Default, Debug)] +pub struct ExtensionData { + entries: Mutex>, +} + +impl ExtensionData { + /// Creates an empty attachment map. + pub fn new() -> Self { + Self::default() + } + + /// Returns the attached value of type `T`, if one exists. + pub fn get(&self) -> Option> + where + T: Any + Send + Sync, + { + let value = self.entries().get(&TypeId::of::())?.clone(); + Some(downcast_data(value)) + } + + /// Returns the attached value of type `T`, inserting one from `init` when absent. + /// + /// The initializer runs while this map is locked, so it should stay cheap; + /// heavyweight lazy work belongs inside the attached value itself. + pub fn get_or_init(&self, init: impl FnOnce() -> T) -> Arc + where + T: Any + Send + Sync, + { + let mut entries = self.entries(); + let value = entries + .entry(TypeId::of::()) + .or_insert_with(|| Arc::new(init())); + downcast_data(Arc::clone(value)) + } + + /// Stores `value` as the attachment of type `T`, returning any previous value. + pub fn insert(&self, value: T) -> Option> + where + T: Any + Send + Sync, + { + self.entries() + .insert(TypeId::of::(), Arc::new(value)) + .map(downcast_data) + } + + /// Removes and returns the attached value of type `T`, if one exists. + pub fn remove(&self) -> Option> + where + T: Any + Send + Sync, + { + self.entries().remove(&TypeId::of::()).map(downcast_data) + } + + fn entries(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.entries.lock().unwrap_or_else(PoisonError::into_inner) + } +} + +fn downcast_data(value: ErasedData) -> Arc +where + T: Any + Send + Sync, +{ + let Ok(value) = value.downcast::() else { + unreachable!("typed extension data stored an incompatible value"); + }; + value +}