Use boxed futures in automation extension

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

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 }

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);
+        })
     }
 }
 

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