feat: async turn item process

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

@@ -1,28 +1,26 @@
 use super::*;
 use codex_extension_api::ExtensionData;
-use codex_extension_api::TurnItemContributionFuture;
 use codex_extension_api::TurnItemContributor;
 use codex_protocol::items::AgentMessageContent;
 use pretty_assertions::assert_eq;
 use std::sync::Arc;
 
 struct RewriteAgentMessageContributor;
 
+#[async_trait::async_trait]
 impl TurnItemContributor for RewriteAgentMessageContributor {
-    fn contribute<'a>(
-        &'a self,
-        _thread_store: &'a ExtensionData,
-        _turn_store: &'a ExtensionData,
-        item: &'a mut TurnItem,
-    ) -> TurnItemContributionFuture<'a> {
-        Box::pin(async move {
-            if let TurnItem::AgentMessage(agent_message) = item {
-                agent_message.content = vec![AgentMessageContent::Text {
-                    text: "plan contributed assistant text".to_string(),
-                }];
-            }
-            Ok(())
-        })
+    async fn contribute(
+        &self,
+        _thread_store: &ExtensionData,
+        _turn_store: &ExtensionData,
+        item: &mut TurnItem,
+    ) -> Result<(), String> {
+        if let TurnItem::AgentMessage(agent_message) = item {
+            agent_message.content = vec![AgentMessageContent::Text {
+                text: "plan contributed assistant text".to_string(),
+            }];
+        }
+        Ok(())
     }
 }
 

codex-rs/core/src/stream_events_utils_tests.rs

@@ -13,7 +13,6 @@ use crate::tools::ToolRouter;
 use crate::tools::parallel::ToolCallRuntime;
 use crate::turn_diff_tracker::TurnDiffTracker;
 use codex_extension_api::ExtensionData;
-use codex_extension_api::TurnItemContributionFuture;
 use codex_extension_api::TurnItemContributor;
 use codex_protocol::error::CodexErr;
 use codex_protocol::items::AgentMessageContent;
@@ -168,43 +167,41 @@ struct TestTurnItemContributor;
 #[derive(Debug)]
 struct TurnItemContributorRan;
 
+#[async_trait::async_trait]
 impl TurnItemContributor for TestTurnItemContributor {
-    fn contribute<'a>(
-        &'a self,
-        _thread_store: &'a ExtensionData,
-        turn_store: &'a ExtensionData,
-        item: &'a mut TurnItem,
-    ) -> TurnItemContributionFuture<'a> {
-        Box::pin(async move {
-            turn_store.insert(TurnItemContributorRan);
-            if let TurnItem::AgentMessage(agent_message) = item {
-                agent_message.memory_citation = Some(MemoryCitation {
-                    entries: Vec::new(),
-                    rollout_ids: Vec::new(),
-                });
-            }
-            Ok(())
-        })
+    async fn contribute(
+        &self,
+        _thread_store: &ExtensionData,
+        turn_store: &ExtensionData,
+        item: &mut TurnItem,
+    ) -> Result<(), String> {
+        turn_store.insert(TurnItemContributorRan);
+        if let TurnItem::AgentMessage(agent_message) = item {
+            agent_message.memory_citation = Some(MemoryCitation {
+                entries: Vec::new(),
+                rollout_ids: Vec::new(),
+            });
+        }
+        Ok(())
     }
 }
 
 struct RewriteAgentMessageContributor;
 
+#[async_trait::async_trait]
 impl TurnItemContributor for RewriteAgentMessageContributor {
-    fn contribute<'a>(
-        &'a self,
-        _thread_store: &'a ExtensionData,
-        _turn_store: &'a ExtensionData,
-        item: &'a mut TurnItem,
-    ) -> TurnItemContributionFuture<'a> {
-        Box::pin(async move {
-            if let TurnItem::AgentMessage(agent_message) = item {
-                agent_message.content = vec![AgentMessageContent::Text {
-                    text: "contributed assistant text".to_string(),
-                }];
-            }
-            Ok(())
-        })
+    async fn contribute(
+        &self,
+        _thread_store: &ExtensionData,
+        _turn_store: &ExtensionData,
+        item: &mut TurnItem,
+    ) -> Result<(), String> {
+        if let TurnItem::AgentMessage(agent_message) = item {
+            agent_message.content = vec![AgentMessageContent::Text {
+                text: "contributed assistant text".to_string(),
+            }];
+        }
+        Ok(())
     }
 }
 

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

@@ -148,20 +148,17 @@ pub trait ApprovalReviewContributor: Send + Sync {
     ) -> Option<ApprovalReviewFuture<'a>>;
 }
 
-/// 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.
+#[async_trait::async_trait]
 pub trait TurnItemContributor: Send + Sync {
-    fn contribute<'a>(
-        &'a self,
-        thread_store: &'a ExtensionData,
-        turn_store: &'a ExtensionData,
-        item: &'a mut TurnItem,
-    ) -> TurnItemContributionFuture<'a>;
+    async fn contribute(
+        &self,
+        thread_store: &ExtensionData,
+        turn_store: &ExtensionData,
+        item: &mut TurnItem,
+    ) -> Result<(), String>;
 }

codex-rs/ext/extension-api/src/lib.rs

@@ -36,7 +36,6 @@ pub use contributors::ToolLifecycleContributor;
 pub use contributors::ToolLifecycleFuture;
 pub use contributors::ToolStartInput;
 pub use contributors::TurnAbortInput;
-pub use contributors::TurnItemContributionFuture;
 pub use contributors::TurnItemContributor;
 pub use contributors::TurnLifecycleContributor;
 pub use contributors::TurnStartInput;