From 905be0fb1fedc66f524c3927105389eb6e208668 Mon Sep 17 00:00:00 2001
From: Andrei Eternal <eternal@openai.com>
Date: Tue, 16 Jun 2026 14:07:02 -0700
Subject: [PATCH] Use boxed futures in automation extension

---
 codex-rs/Cargo.lock                      |  1 -
 codex-rs/ext/automation/Cargo.toml       |  1 -
 codex-rs/ext/automation/src/extension.rs | 18 +++++++++++-------
 codex-rs/ext/automation/src/tool.rs      | 23 ++++++++++++-----------
 4 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock
index e24a2a8130..8e2ff1c846 100644
--- a/codex-rs/Cargo.lock
+++ b/codex-rs/Cargo.lock
@@ -2240,7 +2240,6 @@ name = "codex-automation-extension"
 version = "0.0.0"
 dependencies = [
  "anyhow",
- "async-trait",
  "chrono",
  "codex-core",
  "codex-extension-api",
diff --git a/codex-rs/ext/automation/Cargo.toml b/codex-rs/ext/automation/Cargo.toml
index 038d0358ea..522ef42bb2 100644
--- a/codex-rs/ext/automation/Cargo.toml
+++ b/codex-rs/ext/automation/Cargo.toml
@@ -14,7 +14,6 @@ doctest = false
 workspace = true
 
 [dependencies]
-async-trait = { workspace = true }
 chrono = { workspace = true }
 codex-core = { workspace = true }
 codex-extension-api = { workspace = true }
diff --git a/codex-rs/ext/automation/src/extension.rs b/codex-rs/ext/automation/src/extension.rs
index b8df23bcea..3ca9b1bf7b 100644
--- a/codex-rs/ext/automation/src/extension.rs
+++ b/codex-rs/ext/automation/src/extension.rs
@@ -1,10 +1,10 @@
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use async_trait::async_trait;
 use codex_core::config::Config;
 use codex_extension_api::ConfigContributor;
 use codex_extension_api::ExtensionData;
+use codex_extension_api::ExtensionFuture;
 use codex_extension_api::ExtensionRegistryBuilder;
 use codex_extension_api::ThreadLifecycleContributor;
 use codex_extension_api::ThreadStartInput;
@@ -95,13 +95,17 @@ impl std::fmt::Debug for AutomationExtension {
     }
 }
 
-#[async_trait]
 impl ThreadLifecycleContributor<Config> for AutomationExtension {
-    async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) {
-        let Some(context) = AutomationThreadContext::from_thread_start(&input) else {
-            return;
-        };
-        input.thread_store.insert(context);
+    fn on_thread_start<'a>(
+        &'a self,
+        input: ThreadStartInput<'a, Config>,
+    ) -> ExtensionFuture<'a, ()> {
+        Box::pin(async move {
+            let Some(context) = AutomationThreadContext::from_thread_start(&input) else {
+                return;
+            };
+            input.thread_store.insert(context);
+        })
     }
 }
 
diff --git a/codex-rs/ext/automation/src/tool.rs b/codex-rs/ext/automation/src/tool.rs
index 7cccbdb096..0edc10fd83 100644
--- a/codex-rs/ext/automation/src/tool.rs
+++ b/codex-rs/ext/automation/src/tool.rs
@@ -1,11 +1,11 @@
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use async_trait::async_trait;
 use codex_extension_api::FunctionCallError;
 use codex_extension_api::JsonToolOutput;
 use codex_extension_api::ToolCall;
 use codex_extension_api::ToolExecutor;
+use codex_extension_api::ToolExecutorFuture;
 use codex_extension_api::ToolName;
 use codex_extension_api::ToolOutput;
 use codex_extension_api::ToolSpec;
@@ -112,7 +112,6 @@ enum AutomationTargetResponse {
     Heartbeat { thread_id: String },
 }
 
-#[async_trait]
 impl ToolExecutor<ToolCall> for AutomationUpdateTool {
     fn tool_name(&self) -> ToolName {
         ToolName::plain(AUTOMATION_UPDATE_TOOL_NAME)
@@ -122,15 +121,17 @@ impl ToolExecutor<ToolCall> for AutomationUpdateTool {
         create_automation_update_tool()
     }
 
-    async fn handle(&self, invocation: ToolCall) -> Result<Box<dyn ToolOutput>, FunctionCallError> {
-        let request: AutomationToolRequest = parse_arguments(invocation.function_arguments()?)?;
-        match request.mode {
-            AutomationToolMode::List => self.handle_list().await,
-            AutomationToolMode::Read => self.handle_read(request).await,
-            AutomationToolMode::Create => self.handle_create(request).await,
-            AutomationToolMode::Update => self.handle_update(request).await,
-            AutomationToolMode::Delete => self.handle_delete(request).await,
-        }
+    fn handle(&self, invocation: ToolCall) -> ToolExecutorFuture<'_> {
+        Box::pin(async move {
+            let request: AutomationToolRequest = parse_arguments(invocation.function_arguments()?)?;
+            match request.mode {
+                AutomationToolMode::List => self.handle_list().await,
+                AutomationToolMode::Read => self.handle_read(request).await,
+                AutomationToolMode::Create => self.handle_create(request).await,
+                AutomationToolMode::Update => self.handle_update(request).await,
+                AutomationToolMode::Delete => self.handle_delete(request).await,
+            }
+        })
     }
 }