Create an agent graph store interface and implementation with the state db

codex-rs/Cargo.lock

@@ -1748,6 +1748,20 @@ dependencies = [
  "unicode-width 0.2.1",
 ]
 
+[[package]]
+name = "codex-agent-graph-store"
+version = "0.0.0"
+dependencies = [
+ "async-trait",
+ "codex-protocol",
+ "codex-state",
+ "pretty_assertions",
+ "serde",
+ "tempfile",
+ "thiserror 2.0.18",
+ "tokio",
+]
+
 [[package]]
 name = "codex-agent-identity"
 version = "0.0.0"

codex-rs/Cargo.toml

@@ -2,6 +2,7 @@
 members = [
     "aws-auth",
     "analytics",
+    "agent-graph-store",
     "agent-identity",
     "backend-client",
     "ansi-escape",
@@ -113,6 +114,7 @@ license = "Apache-2.0"
 # Internal
 app_test_support = { path = "app-server/tests/common" }
 codex-analytics = { path = "analytics" }
+codex-agent-graph-store = { path = "agent-graph-store" }
 codex-agent-identity = { path = "agent-identity" }
 codex-ansi-escape = { path = "ansi-escape" }
 codex-api = { path = "codex-api" }

codex-rs/agent-graph-store/BUILD.bazel

@@ -0,0 +1,6 @@
+load("//:defs.bzl", "codex_rust_crate")
+
+codex_rust_crate(
+    name = "agent-graph-store",
+    crate_name = "codex_agent_graph_store",
+)

codex-rs/agent-graph-store/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+edition.workspace = true
+license.workspace = true
+name = "codex-agent-graph-store"
+version.workspace = true
+
+[lib]
+name = "codex_agent_graph_store"
+path = "src/lib.rs"
+
+[lints]
+workspace = true
+
+[dependencies]
+async-trait = { workspace = true }
+codex-protocol = { workspace = true }
+codex-state = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+thiserror = { workspace = true }
+
+[dev-dependencies]
+pretty_assertions = { workspace = true }
+tempfile = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

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

@@ -0,0 +1,20 @@
+/// Result type returned by agent graph store operations.
+pub type AgentGraphStoreResult<T> = Result<T, AgentGraphStoreError>;
+
+/// Error type shared by agent graph store implementations.
+#[derive(Debug, thiserror::Error)]
+pub enum AgentGraphStoreError {
+    /// The caller supplied invalid request data.
+    #[error("invalid agent graph store request: {message}")]
+    InvalidRequest {
+        /// User-facing explanation of the invalid request.
+        message: String,
+    },
+
+    /// Catch-all for implementation failures that do not fit a more specific category.
+    #[error("agent graph store internal error: {message}")]
+    Internal {
+        /// User-facing explanation of the implementation failure.
+        message: String,
+    },
+}

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

@@ -0,0 +1,12 @@
+//! Storage-neutral parent/child topology for thread-spawned agents.
+
+mod error;
+mod local;
+mod store;
+mod types;
+
+pub use error::AgentGraphStoreError;
+pub use error::AgentGraphStoreResult;
+pub use local::LocalAgentGraphStore;
+pub use store::AgentGraphStore;
+pub use types::ThreadSpawnEdgeStatus;

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

@@ -0,0 +1,340 @@
+use async_trait::async_trait;
+use codex_protocol::ThreadId;
+use codex_state::StateRuntime;
+use std::sync::Arc;
+
+use crate::AgentGraphStore;
+use crate::AgentGraphStoreError;
+use crate::AgentGraphStoreResult;
+use crate::ThreadSpawnEdgeStatus;
+
+/// SQLite-backed implementation of [`AgentGraphStore`] using an existing state runtime.
+#[derive(Clone)]
+pub struct LocalAgentGraphStore {
+    state_db: Arc<StateRuntime>,
+}
+
+impl std::fmt::Debug for LocalAgentGraphStore {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LocalAgentGraphStore")
+            .field("codex_home", &self.state_db.codex_home())
+            .finish_non_exhaustive()
+    }
+}
+
+impl LocalAgentGraphStore {
+    /// Create a local graph store from an already-initialized state runtime.
+    pub fn new(state_db: Arc<StateRuntime>) -> Self {
+        Self { state_db }
+    }
+}
+
+#[async_trait]
+impl AgentGraphStore for LocalAgentGraphStore {
+    async 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)
+    }
+
+    async 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)
+    }
+
+    async 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);
+        }
+
+        let mut open_children = self
+            .state_db
+            .list_thread_spawn_children_with_status(
+                parent_thread_id,
+                codex_state::DirectionalThreadSpawnEdgeStatus::Open,
+            )
+            .await
+            .map_err(internal_error)?;
+        let closed_children = self
+            .state_db
+            .list_thread_spawn_children_with_status(
+                parent_thread_id,
+                codex_state::DirectionalThreadSpawnEdgeStatus::Closed,
+            )
+            .await
+            .map_err(internal_error)?;
+        open_children.extend(closed_children);
+        open_children.sort_by_key(std::string::ToString::to_string);
+        Ok(open_children)
+    }
+
+    async 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),
+        }
+    }
+}
+
+fn to_state_status(status: ThreadSpawnEdgeStatus) -> codex_state::DirectionalThreadSpawnEdgeStatus {
+    match status {
+        ThreadSpawnEdgeStatus::Open => codex_state::DirectionalThreadSpawnEdgeStatus::Open,
+        ThreadSpawnEdgeStatus::Closed => codex_state::DirectionalThreadSpawnEdgeStatus::Closed,
+    }
+}
+
+fn internal_error(err: impl std::fmt::Display) -> AgentGraphStoreError {
+    AgentGraphStoreError::Internal {
+        message: err.to_string(),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use codex_state::DirectionalThreadSpawnEdgeStatus;
+    use pretty_assertions::assert_eq;
+    use tempfile::TempDir;
+
+    struct TestRuntime {
+        state_db: Arc<StateRuntime>,
+        _codex_home: TempDir,
+    }
+
+    fn thread_id(suffix: u128) -> ThreadId {
+        ThreadId::from_string(&format!("00000000-0000-0000-0000-{suffix:012}"))
+            .expect("valid thread id")
+    }
+
+    async fn state_runtime() -> TestRuntime {
+        let codex_home = TempDir::new().expect("tempdir should be created");
+        let state_db =
+            StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string())
+                .await
+                .expect("state db should initialize");
+        TestRuntime {
+            state_db,
+            _codex_home: codex_home,
+        }
+    }
+
+    #[tokio::test]
+    async fn local_store_upserts_and_lists_direct_children_with_status_filters() {
+        let fixture = state_runtime().await;
+        let state_db = fixture.state_db;
+        let store = LocalAgentGraphStore::new(state_db.clone());
+        let parent_thread_id = thread_id(1);
+        let first_child_thread_id = thread_id(2);
+        let second_child_thread_id = thread_id(3);
+
+        store
+            .upsert_thread_spawn_edge(
+                parent_thread_id,
+                second_child_thread_id,
+                ThreadSpawnEdgeStatus::Closed,
+            )
+            .await
+            .expect("closed child edge should insert");
+        store
+            .upsert_thread_spawn_edge(
+                parent_thread_id,
+                first_child_thread_id,
+                ThreadSpawnEdgeStatus::Open,
+            )
+            .await
+            .expect("open child edge should insert");
+
+        let all_children = store
+            .list_thread_spawn_children(parent_thread_id, None)
+            .await
+            .expect("all children should load");
+        assert_eq!(
+            all_children,
+            vec![first_child_thread_id, second_child_thread_id]
+        );
+
+        let open_children = store
+            .list_thread_spawn_children(parent_thread_id, Some(ThreadSpawnEdgeStatus::Open))
+            .await
+            .expect("open children should load");
+        let state_open_children = state_db
+            .list_thread_spawn_children_with_status(
+                parent_thread_id,
+                DirectionalThreadSpawnEdgeStatus::Open,
+            )
+            .await
+            .expect("state open children should load");
+        assert_eq!(open_children, state_open_children);
+        assert_eq!(open_children, vec![first_child_thread_id]);
+
+        let closed_children = store
+            .list_thread_spawn_children(parent_thread_id, Some(ThreadSpawnEdgeStatus::Closed))
+            .await
+            .expect("closed children should load");
+        assert_eq!(closed_children, vec![second_child_thread_id]);
+    }
+
+    #[tokio::test]
+    async fn local_store_updates_edge_status() {
+        let fixture = state_runtime().await;
+        let state_db = fixture.state_db;
+        let store = LocalAgentGraphStore::new(state_db);
+        let parent_thread_id = thread_id(10);
+        let child_thread_id = thread_id(11);
+
+        store
+            .upsert_thread_spawn_edge(
+                parent_thread_id,
+                child_thread_id,
+                ThreadSpawnEdgeStatus::Open,
+            )
+            .await
+            .expect("child edge should insert");
+        store
+            .set_thread_spawn_edge_status(child_thread_id, ThreadSpawnEdgeStatus::Closed)
+            .await
+            .expect("child edge should close");
+
+        let open_children = store
+            .list_thread_spawn_children(parent_thread_id, Some(ThreadSpawnEdgeStatus::Open))
+            .await
+            .expect("open children should load");
+        assert_eq!(open_children, Vec::<ThreadId>::new());
+
+        let closed_children = store
+            .list_thread_spawn_children(parent_thread_id, Some(ThreadSpawnEdgeStatus::Closed))
+            .await
+            .expect("closed children should load");
+        assert_eq!(closed_children, vec![child_thread_id]);
+    }
+
+    #[tokio::test]
+    async fn local_store_lists_descendants_breadth_first_with_status_filters() {
+        let fixture = state_runtime().await;
+        let state_db = fixture.state_db;
+        let store = LocalAgentGraphStore::new(state_db.clone());
+        let root_thread_id = thread_id(20);
+        let later_child_thread_id = thread_id(22);
+        let earlier_child_thread_id = thread_id(21);
+        let closed_grandchild_thread_id = thread_id(23);
+        let open_grandchild_thread_id = thread_id(24);
+        let closed_child_thread_id = thread_id(25);
+        let closed_great_grandchild_thread_id = thread_id(26);
+
+        for (parent_thread_id, child_thread_id, status) in [
+            (
+                root_thread_id,
+                later_child_thread_id,
+                ThreadSpawnEdgeStatus::Open,
+            ),
+            (
+                root_thread_id,
+                earlier_child_thread_id,
+                ThreadSpawnEdgeStatus::Open,
+            ),
+            (
+                earlier_child_thread_id,
+                open_grandchild_thread_id,
+                ThreadSpawnEdgeStatus::Open,
+            ),
+            (
+                later_child_thread_id,
+                closed_grandchild_thread_id,
+                ThreadSpawnEdgeStatus::Closed,
+            ),
+            (
+                root_thread_id,
+                closed_child_thread_id,
+                ThreadSpawnEdgeStatus::Closed,
+            ),
+            (
+                closed_child_thread_id,
+                closed_great_grandchild_thread_id,
+                ThreadSpawnEdgeStatus::Closed,
+            ),
+        ] {
+            store
+                .upsert_thread_spawn_edge(parent_thread_id, child_thread_id, status)
+                .await
+                .expect("edge should insert");
+        }
+
+        let all_descendants = store
+            .list_thread_spawn_descendants(root_thread_id, None)
+            .await
+            .expect("all descendants should load");
+        assert_eq!(
+            all_descendants,
+            vec![
+                earlier_child_thread_id,
+                later_child_thread_id,
+                closed_child_thread_id,
+                closed_grandchild_thread_id,
+                open_grandchild_thread_id,
+                closed_great_grandchild_thread_id,
+            ]
+        );
+
+        let open_descendants = store
+            .list_thread_spawn_descendants(root_thread_id, Some(ThreadSpawnEdgeStatus::Open))
+            .await
+            .expect("open descendants should load");
+        let state_open_descendants = state_db
+            .list_thread_spawn_descendants_with_status(
+                root_thread_id,
+                DirectionalThreadSpawnEdgeStatus::Open,
+            )
+            .await
+            .expect("state open descendants should load");
+        assert_eq!(open_descendants, state_open_descendants);
+        assert_eq!(
+            open_descendants,
+            vec![
+                earlier_child_thread_id,
+                later_child_thread_id,
+                open_grandchild_thread_id,
+            ]
+        );
+
+        let closed_descendants = store
+            .list_thread_spawn_descendants(root_thread_id, Some(ThreadSpawnEdgeStatus::Closed))
+            .await
+            .expect("closed descendants should load");
+        assert_eq!(
+            closed_descendants,
+            vec![closed_child_thread_id, closed_great_grandchild_thread_id]
+        );
+    }
+}

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

@@ -0,0 +1,38 @@
+use async_trait::async_trait;
+use codex_protocol::ThreadId;
+
+use crate::AgentGraphStoreResult;
+use crate::ThreadSpawnEdgeStatus;
+
+/// Storage-neutral boundary for persisted thread-spawn parent/child topology.
+#[async_trait]
+pub trait AgentGraphStore: Send + Sync {
+    /// Insert or replace the directional parent/child edge for a spawned thread.
+    async fn upsert_thread_spawn_edge(
+        &self,
+        parent_thread_id: ThreadId,
+        child_thread_id: ThreadId,
+        status: ThreadSpawnEdgeStatus,
+    ) -> AgentGraphStoreResult<()>;
+
+    /// Update the persisted lifecycle status of a spawned thread's incoming edge.
+    async fn set_thread_spawn_edge_status(
+        &self,
+        child_thread_id: ThreadId,
+        status: ThreadSpawnEdgeStatus,
+    ) -> AgentGraphStoreResult<()>;
+
+    /// List direct spawned children of a parent thread.
+    async fn list_thread_spawn_children(
+        &self,
+        parent_thread_id: ThreadId,
+        status_filter: Option<ThreadSpawnEdgeStatus>,
+    ) -> AgentGraphStoreResult<Vec<ThreadId>>;
+
+    /// List spawned descendants breadth-first by depth, then by thread id.
+    async fn list_thread_spawn_descendants(
+        &self,
+        root_thread_id: ThreadId,
+        status_filter: Option<ThreadSpawnEdgeStatus>,
+    ) -> AgentGraphStoreResult<Vec<ThreadId>>;
+}

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

@@ -0,0 +1,11 @@
+use serde::Deserialize;
+use serde::Serialize;
+
+/// Lifecycle status attached to a directional thread-spawn edge.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum ThreadSpawnEdgeStatus {
+    /// The child thread is still live or resumable as an open spawned agent.
+    Open,
+    /// The child thread has been closed from the parent/child graph's perspective.
+    Closed,
+}