[codex] Inject agent graph store into ThreadManager (#29736)

Pick up the AgentGraphStore migration.

- Inject an explicit optional agent graph store into `ThreadManager` 
- Move all calls to spawn, close, recursive resume, and
subtree/archive/delete/feedback traversal through it
- Keep using  `LocalAgentGraphStore` when SQLite is available

This required some changes to the interface to deal with futures:

- The interface now matches `ThreadStore`'s object-safe pattern by
returning a boxed `AgentGraphStoreFuture` directly, allowing
`ThreadManager` to hold `Arc<dyn AgentGraphStore>`

*Slight behavior change!* Unfiltered subtree enumeration now performs a
single all-status breadth-first traversal, so a closed grandchild
beneath an open edge is included; the previous Open-then-Closed
traversals could not cross mixed-status paths and silently omitted it.

codex-rs/agent-graph-store/src/lib.rs

@@ -9,4 +9,5 @@ pub use error::AgentGraphStoreError;
 pub use error::AgentGraphStoreResult;
 pub use local::LocalAgentGraphStore;
 pub use store::AgentGraphStore;
+pub use store::AgentGraphStoreFuture;
 pub use types::ThreadSpawnEdgeStatus;

codex-rs/agent-graph-store/src/local.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use crate::AgentGraphStore;
 use crate::AgentGraphStoreError;
-use crate::AgentGraphStoreResult;
+use crate::AgentGraphStoreFuture;
 use crate::ThreadSpawnEdgeStatus;
 
 /// SQLite-backed implementation of [`AgentGraphStore`] using an existing state runtime.
@@ -29,65 +29,83 @@ impl LocalAgentGraphStore {
 }
 
 impl AgentGraphStore for LocalAgentGraphStore {
-    async fn upsert_thread_spawn_edge(
+    fn upsert_thread_spawn_edge(
         &self,
         parent_thread_id: ThreadId,
         child_thread_id: ThreadId,
         status: ThreadSpawnEdgeStatus,
-    ) -> AgentGraphStoreResult<()> {
-        self.state_db
-            .upsert_thread_spawn_edge(parent_thread_id, child_thread_id, to_state_status(status))
-            .await
-            .map_err(internal_error)
+    ) -> AgentGraphStoreFuture<'_, ()> {
+        Box::pin(async move {
+            self.state_db
+                .upsert_thread_spawn_edge(
+                    parent_thread_id,
+                    child_thread_id,
+                    to_state_status(status),
+                )
+                .await
+                .map_err(internal_error)
+        })
     }
 
-    async fn set_thread_spawn_edge_status(
+    fn set_thread_spawn_edge_status(
         &self,
         child_thread_id: ThreadId,
         status: ThreadSpawnEdgeStatus,
-    ) -> AgentGraphStoreResult<()> {
-        self.state_db
-            .set_thread_spawn_edge_status(child_thread_id, to_state_status(status))
-            .await
-            .map_err(internal_error)
+    ) -> AgentGraphStoreFuture<'_, ()> {
+        Box::pin(async move {
+            self.state_db
+                .set_thread_spawn_edge_status(child_thread_id, to_state_status(status))
+                .await
+                .map_err(internal_error)
+        })
     }
 
-    async fn list_thread_spawn_children(
+    fn list_thread_spawn_children(
         &self,
         parent_thread_id: ThreadId,
         status_filter: Option<ThreadSpawnEdgeStatus>,
-    ) -> AgentGraphStoreResult<Vec<ThreadId>> {
-        if let Some(status) = status_filter {
-            return self
-                .state_db
-                .list_thread_spawn_children_with_status(parent_thread_id, to_state_status(status))
-                .await
-                .map_err(internal_error);
-        }
+    ) -> AgentGraphStoreFuture<'_, Vec<ThreadId>> {
+        Box::pin(async move {
+            if let Some(status) = status_filter {
+                return self
+                    .state_db
+                    .list_thread_spawn_children_with_status(
+                        parent_thread_id,
+                        to_state_status(status),
+                    )
+                    .await
+                    .map_err(internal_error);
+            }
 
-        self.state_db
-            .list_thread_spawn_children(parent_thread_id)
-            .await
-            .map_err(internal_error)
+            self.state_db
+                .list_thread_spawn_children(parent_thread_id)
+                .await
+                .map_err(internal_error)
+        })
     }
 
-    async fn list_thread_spawn_descendants(
+    fn list_thread_spawn_descendants(
         &self,
         root_thread_id: ThreadId,
         status_filter: Option<ThreadSpawnEdgeStatus>,
-    ) -> AgentGraphStoreResult<Vec<ThreadId>> {
-        match status_filter {
-            Some(status) => self
-                .state_db
-                .list_thread_spawn_descendants_with_status(root_thread_id, to_state_status(status))
-                .await
-                .map_err(internal_error),
-            None => self
-                .state_db
-                .list_thread_spawn_descendants(root_thread_id)
-                .await
-                .map_err(internal_error),
-        }
+    ) -> AgentGraphStoreFuture<'_, Vec<ThreadId>> {
+        Box::pin(async move {
+            match status_filter {
+                Some(status) => self
+                    .state_db
+                    .list_thread_spawn_descendants_with_status(
+                        root_thread_id,
+                        to_state_status(status),
+                    )
+                    .await
+                    .map_err(internal_error),
+                None => self
+                    .state_db
+                    .list_thread_spawn_descendants(root_thread_id)
+                    .await
+                    .map_err(internal_error),
+            }
+        })
     }
 }
 

codex-rs/agent-graph-store/src/store.rs

@@ -1,8 +1,15 @@
+use std::future::Future;
+use std::pin::Pin;
+
 use codex_protocol::ThreadId;
 
 use crate::AgentGraphStoreResult;
 use crate::ThreadSpawnEdgeStatus;
 
+/// Future returned by [`AgentGraphStore`] operations.
+pub type AgentGraphStoreFuture<'a, T> =
+    Pin<Box<dyn Future<Output = AgentGraphStoreResult<T>> + Send + 'a>>;
+
 /// Storage-neutral boundary for persisted thread-spawn parent/child topology.
 ///
 /// Implementations are expected to return stable ordering for list methods so callers can merge
@@ -17,7 +24,7 @@ pub trait AgentGraphStore: Send + Sync {
         parent_thread_id: ThreadId,
         child_thread_id: ThreadId,
         status: ThreadSpawnEdgeStatus,
-    ) -> impl std::future::Future<Output = AgentGraphStoreResult<()>> + Send;
+    ) -> AgentGraphStoreFuture<'_, ()>;
 
     /// Update the persisted lifecycle status of a spawned thread's incoming edge.
     ///
@@ -26,7 +33,7 @@ pub trait AgentGraphStore: Send + Sync {
         &self,
         child_thread_id: ThreadId,
         status: ThreadSpawnEdgeStatus,
-    ) -> impl std::future::Future<Output = AgentGraphStoreResult<()>> + Send;
+    ) -> AgentGraphStoreFuture<'_, ()>;
 
     /// List direct spawned children of a parent thread.
     ///
@@ -37,7 +44,7 @@ pub trait AgentGraphStore: Send + Sync {
         &self,
         parent_thread_id: ThreadId,
         status_filter: Option<ThreadSpawnEdgeStatus>,
-    ) -> impl std::future::Future<Output = AgentGraphStoreResult<Vec<ThreadId>>> + Send;
+    ) -> AgentGraphStoreFuture<'_, Vec<ThreadId>>;
 
     /// List spawned descendants breadth-first by depth, then by thread id.
     ///
@@ -49,5 +56,5 @@ pub trait AgentGraphStore: Send + Sync {
         &self,
         root_thread_id: ThreadId,
         status_filter: Option<ThreadSpawnEdgeStatus>,
-    ) -> impl std::future::Future<Output = AgentGraphStoreResult<Vec<ThreadId>>> + Send;
+    ) -> AgentGraphStoreFuture<'_, Vec<ThreadId>>;
 }