Make response item injection thread scoped

codex-rs/app-server/src/extensions.rs

@@ -13,13 +13,10 @@ use codex_extension_api::ExtensionEventFuture;
 use codex_extension_api::ExtensionEventSink;
 use codex_extension_api::ExtensionRegistry;
 use codex_extension_api::ExtensionRegistryBuilder;
-use codex_extension_api::ResponseItemInjectionFuture;
-use codex_extension_api::ResponseItemInjector;
 use codex_features::Feature;
 use codex_login::AuthManager;
 use codex_protocol::ThreadId;
 use codex_protocol::error::CodexErr;
-use codex_protocol::models::ResponseInputItem;
 use codex_protocol::protocol::Event;
 use codex_protocol::protocol::EventMsg;
 use codex_rollout::state_db::StateDbHandle;
@@ -31,7 +28,6 @@ pub(crate) fn thread_extensions<S>(
     event_sink: Arc<dyn ExtensionEventSink>,
     auth_manager: Arc<AuthManager>,
     state_db: Option<StateDbHandle>,
-    thread_manager: Weak<ThreadManager>,
 ) -> Arc<ExtensionRegistry<Config>>
 where
     S: AgentSpawner<StartThreadOptions, Spawned = NewThread, Error = CodexErr> + 'static,
@@ -43,7 +39,6 @@ where
             &mut builder,
             state_db,
             codex_otel::global(),
-            Arc::new(ThreadManagerResponseItemInjector { thread_manager }),
             |config: &Config| config.features.enabled(Feature::Goals),
         );
     }
@@ -52,29 +47,6 @@ where
     Arc::new(builder.build())
 }
 
-struct ThreadManagerResponseItemInjector {
-    thread_manager: Weak<ThreadManager>,
-}
-
-impl ResponseItemInjector for ThreadManagerResponseItemInjector {
-    fn inject_response_items<'a>(
-        &'a self,
-        thread_id: ThreadId,
-        items: Vec<ResponseInputItem>,
-    ) -> ResponseItemInjectionFuture<'a> {
-        let thread_manager = self.thread_manager.clone();
-        Box::pin(async move {
-            let Some(thread_manager) = thread_manager.upgrade() else {
-                return Err(items);
-            };
-            let Ok(thread) = thread_manager.get_thread(thread_id).await else {
-                return Err(items);
-            };
-            thread.inject_response_items_into_active_turn(items).await
-        })
-    }
-}
-
 pub(crate) fn app_server_extension_event_sink(
     outgoing: Arc<OutgoingMessageSender>,
 ) -> Arc<dyn ExtensionEventSink> {

codex-rs/app-server/src/mcp_refresh.rs

@@ -192,7 +192,6 @@ mod tests {
                     Arc::new(NoopExtensionEventSink),
                     auth_manager.clone(),
                     Some(state_db.clone()),
-                    thread_manager.clone(),
                 ),
                 /*analytics_events_client*/ None,
                 Arc::clone(&thread_store),

codex-rs/app-server/src/message_processor.rs

@@ -316,7 +316,6 @@ impl MessageProcessor {
                     app_server_extension_event_sink(outgoing.clone()),
                     auth_manager.clone(),
                     state_db.clone(),
-                    thread_manager.clone(),
                 ),
                 Some(analytics_events_client.clone()),
                 Arc::clone(&thread_store),

codex-rs/core/src/session/session.rs

@@ -10,6 +10,7 @@ use codex_protocol::permissions::FileSystemPath;
 use codex_protocol::permissions::FileSystemSpecialPath;
 use codex_protocol::protocol::ThreadSource;
 use codex_protocol::protocol::TurnEnvironmentSelection;
+use std::sync::Weak;
 use tokio::sync::Semaphore;
 
 /// Context for an initialized model agent
@@ -37,6 +38,25 @@ pub(crate) struct Session {
     pub(super) next_internal_sub_id: AtomicU64,
 }
 
+struct SessionResponseItemInjector {
+    session: Weak<Session>,
+}
+
+impl codex_extension_api::ResponseItemInjector for SessionResponseItemInjector {
+    fn inject_response_items<'a>(
+        &'a self,
+        items: Vec<ResponseInputItem>,
+    ) -> codex_extension_api::ResponseItemInjectionFuture<'a> {
+        let session = self.session.clone();
+        Box::pin(async move {
+            let Some(session) = session.upgrade() else {
+                return Err(items);
+            };
+            session.inject_response_items(items).await
+        })
+    }
+}
+
 #[derive(Clone)]
 pub(crate) struct SessionConfiguration {
     /// Provider identifier ("openai", "openrouter", ...).
@@ -961,13 +981,6 @@ impl Session {
                 codex_extension_api::ExtensionData::new(session_id.to_string());
             let thread_extension_data =
                 codex_extension_api::ExtensionData::new(thread_id.to_string());
-            for contributor in extensions.thread_lifecycle_contributors() {
-                contributor.on_thread_start(codex_extension_api::ThreadStartInput {
-                    config: config.as_ref(),
-                    session_store: &session_extension_data,
-                    thread_store: &thread_extension_data,
-                }).await;
-            }
 
             let services = SessionServices {
                 // Initialize the MCP connection manager with an uninitialized
@@ -1063,6 +1076,20 @@ impl Session {
                 let mut guard = network_policy_decider_session.write().await;
                 *guard = Arc::downgrade(&sess);
             }
+            let response_item_injector: Arc<dyn codex_extension_api::ResponseItemInjector> =
+                Arc::new(SessionResponseItemInjector {
+                    session: Arc::downgrade(&sess),
+                });
+            for contributor in sess.services.extensions.thread_lifecycle_contributors() {
+                contributor
+                    .on_thread_start(codex_extension_api::ThreadStartInput {
+                        config: config.as_ref(),
+                        response_item_injector: Arc::clone(&response_item_injector),
+                        session_store: &sess.services.session_extension_data,
+                        thread_store: &sess.services.thread_extension_data,
+                    })
+                    .await;
+            }
             // Dispatch the SessionConfiguredEvent first and then report any errors.
             // If resuming, include converted initial messages in the payload so UIs can render them immediately.
             let initial_messages = initial_history.get_event_msgs();

codex-rs/ext/extension-api/src/capabilities/response_items.rs

@@ -1,22 +1,20 @@
 use std::future::Future;
 use std::pin::Pin;
 
-use codex_protocol::ThreadId;
 use codex_protocol::models::ResponseInputItem;
 
 /// Future returned when an extension asks the host to inject model-visible input.
 pub type ResponseItemInjectionFuture<'a> =
     Pin<Box<dyn Future<Output = Result<(), Vec<ResponseInputItem>>> + Send + 'a>>;
 
-/// Host-provided helper for extensions that need to steer the active model turn.
+/// Thread-scoped host helper for extensions that need to steer the active model turn.
 ///
 /// Implementations should inject the supplied response items into the active turn
-/// when one can accept same-turn model input. If injection is unavailable, they
-/// return the unchanged items to the caller.
+/// for the current thread when it can accept same-turn model input. If injection
+/// is unavailable, they return the unchanged items to the caller.
 pub trait ResponseItemInjector: Send + Sync {
     fn inject_response_items<'a>(
         &'a self,
-        thread_id: ThreadId,
         items: Vec<ResponseInputItem>,
     ) -> ResponseItemInjectionFuture<'a>;
 }
@@ -28,7 +26,6 @@ pub struct NoopResponseItemInjector;
 impl ResponseItemInjector for NoopResponseItemInjector {
     fn inject_response_items<'a>(
         &'a self,
-        _thread_id: ThreadId,
         items: Vec<ResponseInputItem>,
     ) -> ResponseItemInjectionFuture<'a> {
         Box::pin(std::future::ready(Err(items)))

codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs

@@ -1,11 +1,17 @@
+use std::sync::Arc;
+
 use crate::ExtensionData;
 use crate::HiddenContextMarker;
+use crate::ResponseItemInjector;
 use codex_protocol::protocol::ThreadSettingsSnapshot;
 
 /// Input supplied when the host starts a runtime for a thread.
 pub struct ThreadStartInput<'a, C> {
     /// Host configuration visible at thread start.
     pub config: &'a C,
+    /// Host-provided helper for injecting model-visible input into this
+    /// thread's active turn.
+    pub response_item_injector: Arc<dyn ResponseItemInjector>,
     /// Store scoped to the host session runtime.
     pub session_store: &'a ExtensionData,
     /// Store scoped to this thread runtime.

codex-rs/ext/goal/src/extension.rs

@@ -5,7 +5,6 @@ use codex_extension_api::ConfigContributor;
 use codex_extension_api::ExtensionData;
 use codex_extension_api::ExtensionEventSink;
 use codex_extension_api::ExtensionRegistryBuilder;
-use codex_extension_api::ResponseItemInjector;
 use codex_extension_api::ThreadIdleInput;
 use codex_extension_api::ThreadIdleTurnContributor;
 use codex_extension_api::ThreadIdleTurnRequestFuture;
@@ -48,7 +47,6 @@ pub struct GoalExtension<C> {
     state_dbs: Arc<codex_state::StateRuntime>,
     event_emitter: GoalEventEmitter,
     metrics: GoalMetrics,
-    response_item_injector: Arc<dyn ResponseItemInjector>,
     goals_enabled: Arc<dyn Fn(&C) -> bool + Send + Sync>,
 }
 
@@ -63,14 +61,12 @@ impl<C> GoalExtension<C> {
         state_dbs: Arc<codex_state::StateRuntime>,
         event_sink: Arc<dyn ExtensionEventSink>,
         metrics_client: Option<MetricsClient>,
-        response_item_injector: Arc<dyn ResponseItemInjector>,
         goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
     ) -> Self {
         Self {
             state_dbs,
             event_emitter: GoalEventEmitter::new(event_sink),
             metrics: GoalMetrics::new(metrics_client),
-            response_item_injector,
             goals_enabled: Arc::new(goals_enabled),
         }
     }
@@ -95,7 +91,7 @@ where
                 Arc::clone(&self.state_dbs),
                 self.event_emitter.clone(),
                 self.metrics.clone(),
-                Arc::clone(&self.response_item_injector),
+                Arc::clone(&input.response_item_injector),
                 accounting_state,
                 enabled,
             )
@@ -455,7 +451,6 @@ pub fn install_with_backend<C>(
     registry: &mut ExtensionRegistryBuilder<C>,
     state_dbs: Arc<codex_state::StateRuntime>,
     metrics_client: Option<MetricsClient>,
-    response_item_injector: Arc<dyn ResponseItemInjector>,
     goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static,
 ) where
     C: Send + Sync + 'static,
@@ -464,7 +459,6 @@ pub fn install_with_backend<C>(
         state_dbs,
         registry.event_sink(),
         metrics_client,
-        response_item_injector,
         goals_enabled,
     ));
     registry.thread_lifecycle_contributor(extension.clone());

codex-rs/ext/goal/src/runtime.rs

@@ -327,7 +327,7 @@ impl GoalRuntimeHandle {
         if self
             .inner
             .response_item_injector
-            .inject_response_items(self.inner.thread_id, vec![item])
+            .inject_response_items(vec![item])
             .await
             .is_err()
         {

codex-rs/ext/goal/tests/goal_extension_backend.rs

@@ -1112,7 +1112,6 @@ async fn installed_tools(
         &mut builder,
         runtime,
         /*metrics_client*/ None,
-        Arc::new(NoopResponseItemInjector),
         |_| true,
     );
     let registry = builder.build();
@@ -1122,6 +1121,7 @@ async fn installed_tools(
         contributor
             .on_thread_start(ThreadStartInput {
                 config: &(),
+                response_item_injector: Arc::new(NoopResponseItemInjector),
                 session_store: &session_store,
                 thread_store: &thread_store,
             })
@@ -1161,7 +1161,6 @@ impl GoalExtensionHarness {
             &mut builder,
             runtime,
             /*metrics_client*/ None,
-            Arc::new(NoopResponseItemInjector),
             |_| true,
         );
         let registry = builder.build();
@@ -1171,6 +1170,7 @@ impl GoalExtensionHarness {
             contributor
                 .on_thread_start(ThreadStartInput {
                     config: &(),
+                    response_item_injector: Arc::new(NoopResponseItemInjector),
                     session_store: &session_store,
                     thread_store: &thread_store,
                 })

codex-rs/mcp-server/src/message_processor.rs

@@ -1,6 +1,5 @@
 use std::collections::HashMap;
 use std::sync::Arc;
-use std::sync::Weak;
 
 use codex_arg0::Arg0DispatchPaths;
 use codex_core::StateDbHandle;
@@ -8,14 +7,11 @@ use codex_core::ThreadManager;
 use codex_core::config::Config;
 use codex_exec_server::EnvironmentManager;
 use codex_extension_api::ExtensionRegistryBuilder;
-use codex_extension_api::ResponseItemInjectionFuture;
-use codex_extension_api::ResponseItemInjector;
 use codex_features::Feature;
 use codex_login::AuthManager;
 use codex_login::default_client::USER_AGENT_SUFFIX;
 use codex_login::default_client::get_codex_user_agent;
 use codex_protocol::ThreadId;
-use codex_protocol::models::ResponseInputItem;
 use codex_protocol::protocol::SessionSource;
 use codex_protocol::protocol::Submission;
 use rmcp::model::CallToolRequestParams;
@@ -69,32 +65,27 @@ impl MessageProcessor {
         )
         .await;
         let thread_store = codex_core::thread_store_from_config(config.as_ref(), state_db.clone());
-        let thread_manager = Arc::new_cyclic(|thread_manager| {
-            let mut extensions = ExtensionRegistryBuilder::<Config>::new();
-            if let Some(state_db) = state_db.clone() {
-                codex_goal_extension::install_with_backend(
-                    &mut extensions,
-                    state_db,
-                    codex_otel::global(),
-                    Arc::new(ThreadManagerResponseItemInjector {
-                        thread_manager: thread_manager.clone(),
-                    }),
-                    |config: &Config| config.features.enabled(Feature::Goals),
-                );
-            }
-            ThreadManager::new(
-                config.as_ref(),
-                auth_manager.clone(),
-                SessionSource::Mcp,
-                Arc::clone(&environment_manager),
-                Arc::new(extensions.build()),
-                /*analytics_events_client*/ None,
-                Arc::clone(&thread_store),
-                state_db.clone(),
-                installation_id.clone(),
-                /*attestation_provider*/ None,
-            )
-        });
+        let mut extensions = ExtensionRegistryBuilder::<Config>::new();
+        if let Some(state_db) = state_db.clone() {
+            codex_goal_extension::install_with_backend(
+                &mut extensions,
+                state_db,
+                codex_otel::global(),
+                |config: &Config| config.features.enabled(Feature::Goals),
+            );
+        }
+        let thread_manager = Arc::new(ThreadManager::new(
+            config.as_ref(),
+            auth_manager,
+            SessionSource::Mcp,
+            Arc::clone(&environment_manager),
+            Arc::new(extensions.build()),
+            /*analytics_events_client*/ None,
+            Arc::clone(&thread_store),
+            state_db.clone(),
+            installation_id.clone(),
+            /*attestation_provider*/ None,
+        ));
         Self {
             outgoing,
             initialized: false,
@@ -628,26 +619,3 @@ impl MessageProcessor {
         tracing::info!("notifications/initialized");
     }
 }
-
-struct ThreadManagerResponseItemInjector {
-    thread_manager: Weak<ThreadManager>,
-}
-
-impl ResponseItemInjector for ThreadManagerResponseItemInjector {
-    fn inject_response_items<'a>(
-        &'a self,
-        thread_id: ThreadId,
-        items: Vec<ResponseInputItem>,
-    ) -> ResponseItemInjectionFuture<'a> {
-        let thread_manager = self.thread_manager.clone();
-        Box::pin(async move {
-            let Some(thread_manager) = thread_manager.upgrade() else {
-                return Err(items);
-            };
-            let Ok(thread) = thread_manager.get_thread(thread_id).await else {
-                return Err(items);
-            };
-            thread.inject_response_items_into_active_turn(items).await
-        })
-    }
-}