extension: add initial typed extension API (#21736)

## Why

`codex-core` still owns a growing amount of product-specific behavior.
This PR starts the extraction path by introducing a small, typed
first-party extension seam: features can install the contribution
families they actually own, while the host keeps lifecycle and state
ownership instead of pushing a broad service locator into the API.

See the `examples/` for illustration

## Known limitations
* Tool contract definition will be shared with core
* Fragments must be extracted
* Missing some contributors

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"

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",

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",
+)

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 }

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(&registry, &session_store, &first_thread_store);
+    contribute_prompt(&registry, &session_store, &first_thread_store);
+    contribute_prompt(&registry, &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<codex_extension_api::PromptFragment> {
+    registry
+        .context_contributors()
+        .iter()
+        .flat_map(|contributor| contributor.contribute(session_store, thread_store))
+        .collect()
+}

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<Self>, 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<PromptFragment> {
+        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<PromptFragment> {
+        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::<ContributionCounts>()
+        .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::<ContributionCounts>()
+        .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<ContributionCounts> {
+    store.get_or_init::<ContributionCounts>(Default::default)
+}

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

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<C>: 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<PromptFragment>;
+}
+
+/// 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<ToolContribution>;
+}
+
+/// Future returned by one ordered turn-item contribution.
+pub type TurnItemContributionFuture<'a> =
+    std::pin::Pin<Box<dyn Future<Output = Result<(), String>> + 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;
+}

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<String>) -> Self {
+        Self {
+            slot,
+            text: text.into(),
+        }
+    }
+
+    /// Creates a developer-policy prompt fragment.
+    pub fn developer_policy(text: impl Into<String>) -> Self {
+        Self::new(PromptSlot::DeveloperPolicy, text)
+    }
+
+    /// Creates a developer-capabilities prompt fragment.
+    pub fn developer_capability(text: impl Into<String>) -> Self {
+        Self::new(PromptSlot::DeveloperCapabilities, text)
+    }
+
+    /// Creates a separate top-level developer prompt fragment.
+    pub fn separate_developer(text: impl Into<String>) -> 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
+    }
+}

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<dyn ToolHandler>,
+    supports_parallel_tool_calls: bool,
+}
+
+impl ToolContribution {
+    pub fn new(spec: ResponsesApiTool, handler: Arc<dyn ToolHandler>) -> 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<dyn ToolHandler> {
+        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<Box<dyn Future<Output = Result<Value, ToolCallError>> + 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<String>) -> Self {
+        Self {
+            message: message.into(),
+        }
+    }
+}

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<C>: Send + Sync {
+    /// Registers this extension's concrete typed contributions.
+    fn install(self: Arc<Self>, registry: &mut ExtensionRegistryBuilder<C>);
+}

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;

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<C> {
+    thread_start_contributors: Vec<Arc<dyn ThreadStartContributor<C>>>,
+    context_contributors: Vec<Arc<dyn ContextContributor>>,
+    tool_contributors: Vec<Arc<dyn ToolContributor>>,
+    turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
+    approval_interceptor_contributors: Vec<Arc<dyn ApprovalInterceptorContributor>>,
+}
+
+impl<C> Default for ExtensionRegistryBuilder<C> {
+    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<C> ExtensionRegistryBuilder<C> {
+    /// Creates an empty registry builder.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Installs one extension and returns the builder.
+    #[must_use]
+    pub fn with_extension<E>(mut self, extension: Arc<E>) -> Self
+    where
+        E: CodexExtension<C> + 'static,
+    {
+        self.install_extension(extension);
+        self
+    }
+
+    /// Installs one extension into the registry under construction.
+    pub fn install_extension<E>(&mut self, extension: Arc<E>)
+    where
+        E: CodexExtension<C> + 'static,
+    {
+        extension.install(self);
+    }
+
+    /// Registers one approval interceptor contributor.
+    pub fn approval_interceptor_contributor(
+        &mut self,
+        contributor: Arc<dyn ApprovalInterceptorContributor>,
+    ) {
+        self.approval_interceptor_contributors.push(contributor);
+    }
+
+    /// Registers one thread-start contributor.
+    pub fn thread_start_contributor(&mut self, contributor: Arc<dyn ThreadStartContributor<C>>) {
+        self.thread_start_contributors.push(contributor);
+    }
+
+    /// Registers one prompt contributor.
+    pub fn prompt_contributor(&mut self, contributor: Arc<dyn ContextContributor>) {
+        self.context_contributors.push(contributor);
+    }
+
+    /// Registers one native tool contributor.
+    pub fn tool_contributor(&mut self, contributor: Arc<dyn ToolContributor>) {
+        self.tool_contributors.push(contributor);
+    }
+
+    /// Registers one ordered turn-item contributor.
+    pub fn turn_item_contributor(&mut self, contributor: Arc<dyn TurnItemContributor>) {
+        self.turn_item_contributors.push(contributor);
+    }
+
+    /// Finishes construction and returns the immutable registry.
+    pub fn build(self) -> ExtensionRegistry<C> {
+        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<C> {
+    thread_start_contributors: Vec<Arc<dyn ThreadStartContributor<C>>>,
+    context_contributors: Vec<Arc<dyn ContextContributor>>,
+    tool_contributors: Vec<Arc<dyn ToolContributor>>,
+    turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
+    approval_interceptor_contributors: Vec<Arc<dyn ApprovalInterceptorContributor>>,
+}
+
+impl<C> ExtensionRegistry<C> {
+    /// Returns the registered thread-start contributors.
+    pub fn thread_start_contributors(&self) -> &[Arc<dyn ThreadStartContributor<C>>] {
+        &self.thread_start_contributors
+    }
+
+    /// Returns the registered approval interceptor contributors.
+    pub fn approval_interceptor_contributors(&self) -> &[Arc<dyn ApprovalInterceptorContributor>] {
+        &self.approval_interceptor_contributors
+    }
+
+    /// Returns the registered prompt contributors.
+    pub fn context_contributors(&self) -> &[Arc<dyn ContextContributor>] {
+        &self.context_contributors
+    }
+
+    /// Returns the registered native tool contributors.
+    pub fn tool_contributors(&self) -> &[Arc<dyn ToolContributor>] {
+        &self.tool_contributors
+    }
+
+    /// Returns the registered ordered turn-item contributors.
+    pub fn turn_item_contributors(&self) -> &[Arc<dyn TurnItemContributor>] {
+        &self.turn_item_contributors
+    }
+}
+
+/// Creates an empty shared registry for hosts that do not install extensions.
+pub fn empty_extension_registry<C>() -> Arc<ExtensionRegistry<C>> {
+    Arc::new(ExtensionRegistryBuilder::new().build())
+}

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<dyn Any + Send + Sync>;
+
+/// Typed extension-owned data attached to one host object.
+#[derive(Default, Debug)]
+pub struct ExtensionData {
+    entries: Mutex<HashMap<TypeId, ErasedData>>,
+}
+
+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<T>(&self) -> Option<Arc<T>>
+    where
+        T: Any + Send + Sync,
+    {
+        let value = self.entries().get(&TypeId::of::<T>())?.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<T>(&self, init: impl FnOnce() -> T) -> Arc<T>
+    where
+        T: Any + Send + Sync,
+    {
+        let mut entries = self.entries();
+        let value = entries
+            .entry(TypeId::of::<T>())
+            .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<T>(&self, value: T) -> Option<Arc<T>>
+    where
+        T: Any + Send + Sync,
+    {
+        self.entries()
+            .insert(TypeId::of::<T>(), Arc::new(value))
+            .map(downcast_data)
+    }
+
+    /// Removes and returns the attached value of type `T`, if one exists.
+    pub fn remove<T>(&self) -> Option<Arc<T>>
+    where
+        T: Any + Send + Sync,
+    {
+        self.entries().remove(&TypeId::of::<T>()).map(downcast_data)
+    }
+
+    fn entries(&self) -> std::sync::MutexGuard<'_, HashMap<TypeId, ErasedData>> {
+        self.entries.lock().unwrap_or_else(PoisonError::into_inner)
+    }
+}
+
+fn downcast_data<T>(value: ErasedData) -> Arc<T>
+where
+    T: Any + Send + Sync,
+{
+    let Ok(value) = value.downcast::<T>() else {
+        unreachable!("typed extension data stored an incompatible value");
+    };
+    value
+}