[codex] Test extension API contracts (#26835)

## Why

`codex-extension-api` defines contracts shared by extension crates and
their hosts, but it had no direct test suite. Host and feature tests
cover downstream behavior, while regressions in the API crate's own
typed state, registry ordering, and capability adapters could go
unnoticed.

## What

- Add public-surface integration tests for `ExtensionData`, including
concurrent initialization and poison recovery.
- Cover contributor registration order, approval short-circuiting, event
sink retention, no-op response injection, and closure-based agent
spawning.
- Add the test-only dependencies used by the suite.

## Validation

- `just test -p codex-extension-api`
- `just argument-comment-lint -p codex-extension-api`
- `just bazel-lock-check`

codex-rs/Cargo.lock

@@ -2923,6 +2923,8 @@ dependencies = [
  "codex-context-fragments",
  "codex-protocol",
  "codex-tools",
+ "pretty_assertions",
+ "tokio",
 ]
 
 [[package]]

codex-rs/ext/extension-api/Cargo.toml

@@ -18,3 +18,7 @@ async-trait = { workspace = true }
 codex-context-fragments = { workspace = true }
 codex-protocol = { workspace = true }
 codex-tools = { workspace = true }
+
+[dev-dependencies]
+pretty_assertions = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

codex-rs/ext/extension-api/tests/capabilities.rs

@@ -0,0 +1,56 @@
+use std::sync::Arc;
+use std::sync::Mutex;
+
+use codex_extension_api::AgentSpawnFuture;
+use codex_extension_api::AgentSpawner;
+use codex_extension_api::NoopResponseItemInjector;
+use codex_extension_api::ResponseItemInjector;
+use codex_protocol::ThreadId;
+use codex_protocol::models::ContentItem;
+use codex_protocol::models::ResponseInputItem;
+use pretty_assertions::assert_eq;
+
+#[tokio::test]
+async fn noop_response_item_injector_returns_original_items() {
+    let items = vec![ResponseInputItem::Message {
+        role: "user".to_string(),
+        content: vec![ContentItem::InputText {
+            text: "keep this input".to_string(),
+        }],
+        phase: None,
+    }];
+
+    let returned_items = NoopResponseItemInjector
+        .inject_response_items(items.clone())
+        .await
+        .expect_err("noop injector should reject same-turn injection");
+
+    assert_eq!(returned_items, items);
+}
+
+#[tokio::test]
+async fn closure_agent_spawner_forwards_arguments_and_result() {
+    let calls = Arc::new(Mutex::new(Vec::new()));
+    let recorded_calls = Arc::clone(&calls);
+    let spawner = move |thread_id: ThreadId,
+                        request: String|
+          -> AgentSpawnFuture<'static, usize, &'static str> {
+        recorded_calls
+            .lock()
+            .expect("agent spawn calls lock")
+            .push((thread_id, request.clone()));
+        Box::pin(async move { Ok(request.len()) })
+    };
+    let thread_id =
+        ThreadId::from_string("11111111-1111-4111-8111-111111111111").expect("valid thread id");
+
+    let spawned = spawner
+        .spawn_subagent(thread_id, "delegate this".to_string())
+        .await;
+
+    assert_eq!(spawned, Ok(13));
+    assert_eq!(
+        calls.lock().expect("agent spawn calls lock").as_slice(),
+        [(thread_id, "delegate this".to_string())]
+    );
+}

codex-rs/ext/extension-api/tests/registry.rs

@@ -0,0 +1,370 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+use codex_extension_api::ApprovalReviewContributor;
+use codex_extension_api::ConfigContributor;
+use codex_extension_api::ContextContributor;
+use codex_extension_api::ContextualUserFragment;
+use codex_extension_api::ExtensionData;
+use codex_extension_api::ExtensionEventSink;
+use codex_extension_api::ExtensionRegistryBuilder;
+use codex_extension_api::PromptFragment;
+use codex_extension_api::ThreadLifecycleContributor;
+use codex_extension_api::TokenUsageContributor;
+use codex_extension_api::ToolCall;
+use codex_extension_api::ToolContributor;
+use codex_extension_api::ToolExecutor;
+use codex_extension_api::ToolLifecycleContributor;
+use codex_extension_api::TurnInputContext;
+use codex_extension_api::TurnInputContributor;
+use codex_extension_api::TurnItemContributor;
+use codex_extension_api::TurnLifecycleContributor;
+use codex_extension_api::empty_extension_registry;
+use codex_protocol::items::HookPromptItem;
+use codex_protocol::items::TurnItem;
+use codex_protocol::protocol::Event;
+use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::ReviewDecision;
+use codex_protocol::protocol::WarningEvent;
+use pretty_assertions::assert_eq;
+
+struct AllContributors;
+
+impl ContextContributor for AllContributors {
+    fn contribute<'a>(
+        &'a self,
+        _session_store: &'a ExtensionData,
+        _thread_store: &'a ExtensionData,
+    ) -> Pin<Box<dyn Future<Output = Vec<PromptFragment>> + Send + 'a>> {
+        Box::pin(std::future::ready(Vec::new()))
+    }
+}
+
+#[async_trait::async_trait]
+impl ThreadLifecycleContributor<()> for AllContributors {}
+
+#[async_trait::async_trait]
+impl TurnLifecycleContributor for AllContributors {}
+
+impl ConfigContributor<()> for AllContributors {}
+
+#[async_trait::async_trait]
+impl TokenUsageContributor for AllContributors {}
+
+#[async_trait::async_trait]
+impl TurnInputContributor for AllContributors {
+    async fn contribute(
+        &self,
+        _input: TurnInputContext,
+        _session_store: &ExtensionData,
+        _thread_store: &ExtensionData,
+        _turn_store: &ExtensionData,
+    ) -> Vec<Box<dyn ContextualUserFragment + Send>> {
+        Vec::new()
+    }
+}
+
+impl ToolContributor for AllContributors {
+    fn tools(
+        &self,
+        _session_store: &ExtensionData,
+        _thread_store: &ExtensionData,
+    ) -> Vec<Arc<dyn ToolExecutor<ToolCall>>> {
+        Vec::new()
+    }
+}
+
+impl ToolLifecycleContributor for AllContributors {}
+
+#[async_trait::async_trait]
+impl TurnItemContributor for AllContributors {
+    async fn contribute(
+        &self,
+        _thread_store: &ExtensionData,
+        _turn_store: &ExtensionData,
+        _item: &mut TurnItem,
+    ) -> Result<(), String> {
+        Ok(())
+    }
+}
+
+#[async_trait::async_trait]
+impl ApprovalReviewContributor for AllContributors {
+    async fn contribute(
+        &self,
+        _session_store: &ExtensionData,
+        _thread_store: &ExtensionData,
+        _prompt: &str,
+    ) -> Option<ReviewDecision> {
+        Some(ReviewDecision::ApprovedForSession)
+    }
+}
+
+#[tokio::test]
+async fn build_round_trips_every_contributor_category() {
+    let contributor = Arc::new(AllContributors);
+    let mut builder = ExtensionRegistryBuilder::<()>::new();
+    builder.thread_lifecycle_contributor(contributor.clone());
+    builder.turn_lifecycle_contributor(contributor.clone());
+    builder.config_contributor(contributor.clone());
+    builder.token_usage_contributor(contributor.clone());
+    builder.prompt_contributor(contributor.clone());
+    builder.turn_input_contributor(contributor.clone());
+    builder.tool_contributor(contributor.clone());
+    builder.tool_lifecycle_contributor(contributor.clone());
+    builder.turn_item_contributor(contributor.clone());
+    builder.approval_review_contributor(contributor);
+    let registry = builder.build();
+
+    assert_eq!(registry.thread_lifecycle_contributors().len(), 1);
+    assert_eq!(registry.turn_lifecycle_contributors().len(), 1);
+    assert_eq!(registry.config_contributors().len(), 1);
+    assert_eq!(registry.token_usage_contributors().len(), 1);
+    assert_eq!(registry.context_contributors().len(), 1);
+    assert_eq!(registry.turn_input_contributors().len(), 1);
+    assert_eq!(registry.tool_contributors().len(), 1);
+    assert_eq!(registry.tool_lifecycle_contributors().len(), 1);
+    assert_eq!(registry.turn_item_contributors().len(), 1);
+    assert_eq!(
+        registry
+            .approval_review(
+                &ExtensionData::new("session"),
+                &ExtensionData::new("thread"),
+                "review this",
+            )
+            .await,
+        Some(ReviewDecision::ApprovedForSession)
+    );
+}
+
+struct NamedContextContributor(&'static str);
+
+impl ContextContributor for NamedContextContributor {
+    fn contribute<'a>(
+        &'a self,
+        _session_store: &'a ExtensionData,
+        _thread_store: &'a ExtensionData,
+    ) -> Pin<Box<dyn Future<Output = Vec<PromptFragment>> + Send + 'a>> {
+        Box::pin(std::future::ready(vec![PromptFragment::developer_policy(
+            self.0,
+        )]))
+    }
+}
+
+struct RecordingTurnItemContributor {
+    name: &'static str,
+    calls: Arc<Mutex<Vec<&'static str>>>,
+}
+
+#[async_trait::async_trait]
+impl TurnItemContributor for RecordingTurnItemContributor {
+    async fn contribute(
+        &self,
+        _thread_store: &ExtensionData,
+        _turn_store: &ExtensionData,
+        _item: &mut TurnItem,
+    ) -> Result<(), String> {
+        self.calls
+            .lock()
+            .unwrap_or_else(|error| panic!("turn item calls lock poisoned: {error}"))
+            .push(self.name);
+        Ok(())
+    }
+}
+
+#[tokio::test]
+async fn contributors_preserve_registration_order() {
+    let turn_item_calls = Arc::new(Mutex::new(Vec::new()));
+    let mut builder = ExtensionRegistryBuilder::<()>::new();
+    builder.prompt_contributor(Arc::new(NamedContextContributor("first")));
+    builder.prompt_contributor(Arc::new(NamedContextContributor("second")));
+    for name in ["first", "second"] {
+        builder.turn_item_contributor(Arc::new(RecordingTurnItemContributor {
+            name,
+            calls: Arc::clone(&turn_item_calls),
+        }));
+    }
+    let registry = builder.build();
+    let session_store = ExtensionData::new("session");
+    let thread_store = ExtensionData::new("thread");
+    let turn_store = ExtensionData::new("turn");
+
+    let mut fragments = Vec::new();
+    for contributor in registry.context_contributors() {
+        fragments.extend(contributor.contribute(&session_store, &thread_store).await);
+    }
+    let mut item = TurnItem::HookPrompt(HookPromptItem {
+        id: "item".to_string(),
+        fragments: Vec::new(),
+    });
+    for contributor in registry.turn_item_contributors() {
+        contributor
+            .contribute(&thread_store, &turn_store, &mut item)
+            .await
+            .expect("turn item contribution should succeed");
+    }
+
+    assert_eq!(
+        fragments,
+        vec![
+            PromptFragment::developer_policy("first"),
+            PromptFragment::developer_policy("second"),
+        ]
+    );
+    assert_eq!(
+        turn_item_calls
+            .lock()
+            .expect("turn item calls lock")
+            .as_slice(),
+        ["first", "second"]
+    );
+}
+
+#[derive(Debug, PartialEq, Eq)]
+struct ApprovalCall {
+    contributor: &'static str,
+    session_id: String,
+    thread_id: String,
+    prompt: String,
+}
+
+struct RecordingApprovalContributor {
+    name: &'static str,
+    decision: Option<ReviewDecision>,
+    calls: Arc<Mutex<Vec<ApprovalCall>>>,
+}
+
+#[async_trait::async_trait]
+impl ApprovalReviewContributor for RecordingApprovalContributor {
+    async fn contribute(
+        &self,
+        session_store: &ExtensionData,
+        thread_store: &ExtensionData,
+        prompt: &str,
+    ) -> Option<ReviewDecision> {
+        self.calls
+            .lock()
+            .unwrap_or_else(|error| panic!("approval calls lock poisoned: {error}"))
+            .push(ApprovalCall {
+                contributor: self.name,
+                session_id: session_store.level_id().to_string(),
+                thread_id: thread_store.level_id().to_string(),
+                prompt: prompt.to_string(),
+            });
+        self.decision.clone()
+    }
+}
+
+#[tokio::test]
+async fn approval_review_returns_first_claim_and_short_circuits() {
+    let calls = Arc::new(Mutex::new(Vec::new()));
+    let mut builder = ExtensionRegistryBuilder::<()>::new();
+    for (name, decision) in [
+        ("first", None),
+        ("second", Some(ReviewDecision::Approved)),
+        ("third", Some(ReviewDecision::Denied)),
+    ] {
+        builder.approval_review_contributor(Arc::new(RecordingApprovalContributor {
+            name,
+            decision,
+            calls: Arc::clone(&calls),
+        }));
+    }
+    let registry = builder.build();
+
+    let decision = registry
+        .approval_review(
+            &ExtensionData::new("session-1"),
+            &ExtensionData::new("thread-1"),
+            "allow command?",
+        )
+        .await;
+
+    assert_eq!(decision, Some(ReviewDecision::Approved));
+    assert_eq!(
+        calls.lock().expect("approval calls lock").as_slice(),
+        [
+            ApprovalCall {
+                contributor: "first",
+                session_id: "session-1".to_string(),
+                thread_id: "thread-1".to_string(),
+                prompt: "allow command?".to_string(),
+            },
+            ApprovalCall {
+                contributor: "second",
+                session_id: "session-1".to_string(),
+                thread_id: "thread-1".to_string(),
+                prompt: "allow command?".to_string(),
+            },
+        ]
+    );
+}
+
+#[derive(Default)]
+struct RecordingEventSink {
+    events: Mutex<Vec<(String, String)>>,
+}
+
+impl ExtensionEventSink for RecordingEventSink {
+    fn emit(&self, event: Event) {
+        let EventMsg::Warning(warning) = event.msg else {
+            panic!("test sink only accepts warning events");
+        };
+        self.events
+            .lock()
+            .unwrap_or_else(|error| panic!("recording event sink lock poisoned: {error}"))
+            .push((event.id, warning.message));
+    }
+}
+
+#[test]
+fn custom_event_sink_survives_registry_build() {
+    let sink = Arc::new(RecordingEventSink::default());
+    let builder = ExtensionRegistryBuilder::<()>::with_event_sink(sink.clone());
+
+    builder
+        .event_sink()
+        .emit(warning_event("builder", "before"));
+    let registry = builder.build();
+    registry
+        .event_sink()
+        .emit(warning_event("registry", "after"));
+
+    assert_eq!(
+        sink.events
+            .lock()
+            .expect("recording event sink lock")
+            .as_slice(),
+        [
+            ("builder".to_string(), "before".to_string()),
+            ("registry".to_string(), "after".to_string()),
+        ]
+    );
+}
+
+#[tokio::test]
+async fn empty_registry_does_not_claim_approval_review() {
+    let registry = empty_extension_registry::<()>();
+
+    assert_eq!(
+        registry
+            .approval_review(
+                &ExtensionData::new("session"),
+                &ExtensionData::new("thread"),
+                "unclaimed",
+            )
+            .await,
+        None
+    );
+}
+
+fn warning_event(id: &str, message: &str) -> Event {
+    Event {
+        id: id.to_string(),
+        msg: EventMsg::Warning(WarningEvent {
+            message: message.to_string(),
+        }),
+    }
+}

codex-rs/ext/extension-api/tests/state.rs

@@ -0,0 +1,109 @@
+use std::panic::AssertUnwindSafe;
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use codex_extension_api::ExtensionData;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn typed_values_can_be_inserted_replaced_and_removed() {
+    let data = ExtensionData::new("thread-1");
+
+    assert_eq!(data.insert(/*value*/ 41_u64), None);
+    assert_eq!(data.insert("alpha".to_string()), None);
+    assert_eq!(data.get::<u64>().as_deref(), Some(&41));
+    assert_eq!(
+        data.get::<String>().map(|value| value.as_str().to_string()),
+        Some("alpha".to_string())
+    );
+
+    assert_eq!(data.insert(/*value*/ 42_u64).as_deref(), Some(&41));
+    assert_eq!(data.get::<u64>().as_deref(), Some(&42));
+    assert_eq!(
+        data.remove::<String>()
+            .map(|value| value.as_str().to_string()),
+        Some("alpha".to_string())
+    );
+    assert_eq!(data.get::<String>(), None);
+    assert_eq!(data.get::<u64>().as_deref(), Some(&42));
+}
+
+#[test]
+fn get_or_init_initializes_once_and_returns_shared_value() {
+    const CALLER_COUNT: usize = 8;
+
+    #[derive(Debug, PartialEq, Eq)]
+    struct SharedValue(usize);
+
+    let data = Arc::new(ExtensionData::new("session"));
+    let callers_started = Arc::new(AtomicUsize::new(0));
+    let initialization_count = Arc::new(AtomicUsize::new(0));
+
+    let handles: [_; CALLER_COUNT] = std::array::from_fn(|_| {
+        let data = Arc::clone(&data);
+        let callers_started = Arc::clone(&callers_started);
+        let initialization_count = Arc::clone(&initialization_count);
+        std::thread::spawn(move || {
+            callers_started.fetch_add(1, Ordering::SeqCst);
+            data.get_or_init(|| {
+                initialization_count.fetch_add(1, Ordering::SeqCst);
+                // Keep the first initializer active until every worker has attempted
+                // get_or_init, forcing callers to overlap on the same missing entry.
+                while callers_started.load(Ordering::SeqCst) < CALLER_COUNT {
+                    std::thread::yield_now();
+                }
+                SharedValue(7)
+            })
+        })
+    });
+    let values = handles
+        .into_iter()
+        .map(|handle| handle.join().expect("initializer thread should succeed"))
+        .collect::<Vec<_>>();
+
+    assert_eq!(initialization_count.load(Ordering::SeqCst), 1);
+    assert_eq!(
+        values.iter().map(Arc::as_ref).collect::<Vec<_>>(),
+        vec![&SharedValue(7); CALLER_COUNT]
+    );
+    assert!(
+        values
+            .iter()
+            .skip(1)
+            .all(|value| Arc::ptr_eq(&values[0], value))
+    );
+}
+
+#[test]
+fn stores_are_isolated_and_preserve_level_id() {
+    let session_data = ExtensionData::new("root-1");
+    let thread_data = ExtensionData::new("root-1");
+
+    session_data.insert(/*value*/ 17_u32);
+    thread_data.insert("thread value".to_string());
+
+    assert_eq!(session_data.level_id(), "root-1");
+    assert_eq!(thread_data.level_id(), "root-1");
+    assert_eq!(session_data.get::<u32>().as_deref(), Some(&17));
+    assert_eq!(session_data.get::<String>(), None);
+    assert_eq!(thread_data.get::<u32>(), None);
+    assert_eq!(
+        thread_data
+            .get::<String>()
+            .map(|value| value.as_str().to_string()),
+        Some("thread value".to_string())
+    );
+}
+
+#[test]
+fn store_remains_usable_after_panicking_initializer() {
+    let data = ExtensionData::new("turn-1");
+
+    let result = std::panic::catch_unwind(AssertUnwindSafe(|| {
+        data.get_or_init::<u64>(|| panic!("initializer failed"));
+    }));
+
+    assert!(result.is_err());
+    assert_eq!(*data.get_or_init(|| 99_u64), 99);
+}