diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c8af30147..87b0bd598 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1748,6 +1748,21 @@ 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", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", +] + [[package]] name = "codex-agent-identity" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 217c179be..ad0bdb5f9 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -2,6 +2,7 @@ members = [ "aws-auth", "analytics", + "agent-graph-store", "agent-identity", "backend-client", "ansi-escape", @@ -119,6 +120,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" } @@ -450,6 +452,7 @@ unwrap_used = "deny" # silence the false positive here instead of deleting a real dependency. [workspace.metadata.cargo-shear] ignored = [ + "codex-agent-graph-store", "icu_provider", "openssl-sys", "codex-utils-readiness", diff --git a/codex-rs/agent-graph-store/BUILD.bazel b/codex-rs/agent-graph-store/BUILD.bazel new file mode 100644 index 000000000..96c077e26 --- /dev/null +++ b/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", +) diff --git a/codex-rs/agent-graph-store/Cargo.toml b/codex-rs/agent-graph-store/Cargo.toml new file mode 100644 index 000000000..e221ef61b --- /dev/null +++ b/codex-rs/agent-graph-store/Cargo.toml @@ -0,0 +1,25 @@ +[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 } +serde_json = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/agent-graph-store/src/error.rs b/codex-rs/agent-graph-store/src/error.rs new file mode 100644 index 000000000..ddd8eeef3 --- /dev/null +++ b/codex-rs/agent-graph-store/src/error.rs @@ -0,0 +1,20 @@ +/// Result type returned by agent graph store operations. +pub type AgentGraphStoreResult = Result; + +/// 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, + }, +} diff --git a/codex-rs/agent-graph-store/src/lib.rs b/codex-rs/agent-graph-store/src/lib.rs new file mode 100644 index 000000000..72e8b45e8 --- /dev/null +++ b/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; diff --git a/codex-rs/agent-graph-store/src/local.rs b/codex-rs/agent-graph-store/src/local.rs new file mode 100644 index 000000000..f45874855 --- /dev/null +++ b/codex-rs/agent-graph-store/src/local.rs @@ -0,0 +1,325 @@ +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, +} + +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) -> 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, + ) -> AgentGraphStoreResult> { + 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) + } + + async fn list_thread_spawn_descendants( + &self, + root_thread_id: ThreadId, + status_filter: Option, + ) -> AgentGraphStoreResult> { + 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, + _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(/*suffix*/ 1); + let first_child_thread_id = thread_id(/*suffix*/ 2); + let second_child_thread_id = thread_id(/*suffix*/ 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, /*status_filter*/ 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(/*suffix*/ 10); + let child_thread_id = thread_id(/*suffix*/ 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::::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(/*suffix*/ 20); + let later_child_thread_id = thread_id(/*suffix*/ 22); + let earlier_child_thread_id = thread_id(/*suffix*/ 21); + let closed_grandchild_thread_id = thread_id(/*suffix*/ 23); + let open_grandchild_thread_id = thread_id(/*suffix*/ 24); + let closed_child_thread_id = thread_id(/*suffix*/ 25); + let closed_great_grandchild_thread_id = thread_id(/*suffix*/ 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, /*status_filter*/ 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] + ); + } +} diff --git a/codex-rs/agent-graph-store/src/store.rs b/codex-rs/agent-graph-store/src/store.rs new file mode 100644 index 000000000..c42118211 --- /dev/null +++ b/codex-rs/agent-graph-store/src/store.rs @@ -0,0 +1,55 @@ +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. +/// +/// Implementations are expected to return stable ordering for list methods so callers can merge +/// persisted graph state with live in-memory state without introducing nondeterministic output. +#[async_trait] +pub trait AgentGraphStore: Send + Sync { + /// Insert or replace the directional parent/child edge for a spawned thread. + /// + /// `child_thread_id` has at most one persisted parent. Re-inserting the same child should + /// update both the parent and status to match the supplied values. + 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. + /// + /// Implementations should treat missing children as a successful no-op. + async fn set_thread_spawn_edge_status( + &self, + child_thread_id: ThreadId, + status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreResult<()>; + + /// List direct spawned children of a parent thread. + /// + /// When `status_filter` is `Some`, only child edges with that exact status are returned. When + /// it is `None`, all direct child edges are returned regardless of status, including statuses + /// that may be added by a future store implementation. + async fn list_thread_spawn_children( + &self, + parent_thread_id: ThreadId, + status_filter: Option, + ) -> AgentGraphStoreResult>; + + /// List spawned descendants breadth-first by depth, then by thread id. + /// + /// `status_filter` is applied to every traversed edge, not just to the returned descendants. + /// For example, `Some(Open)` walks only open edges, so descendants under a closed edge are not + /// included even if their own incoming edge is open. `None` walks and returns every persisted + /// edge regardless of status. + async fn list_thread_spawn_descendants( + &self, + root_thread_id: ThreadId, + status_filter: Option, + ) -> AgentGraphStoreResult>; +} diff --git a/codex-rs/agent-graph-store/src/types.rs b/codex-rs/agent-graph-store/src/types.rs new file mode 100644 index 000000000..2a9f6caed --- /dev/null +++ b/codex-rs/agent-graph-store/src/types.rs @@ -0,0 +1,42 @@ +use serde::Deserialize; +use serde::Serialize; + +/// Lifecycle status attached to a directional thread-spawn edge. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +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, +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn thread_spawn_edge_status_serializes_as_snake_case() { + assert_eq!( + serde_json::to_string(&ThreadSpawnEdgeStatus::Open) + .expect("open status should serialize"), + "\"open\"" + ); + assert_eq!( + serde_json::to_string(&ThreadSpawnEdgeStatus::Closed) + .expect("closed status should serialize"), + "\"closed\"" + ); + assert_eq!( + serde_json::from_str::("\"open\"") + .expect("open status should deserialize"), + ThreadSpawnEdgeStatus::Open + ); + assert_eq!( + serde_json::from_str::("\"closed\"") + .expect("closed status should deserialize"), + ThreadSpawnEdgeStatus::Closed + ); + } +} diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 8a3bcee5b..906a3bb39 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -134,6 +134,15 @@ ON CONFLICT(child_thread_id) DO UPDATE SET .await } + /// List all direct spawned children of `parent_thread_id`. + pub async fn list_thread_spawn_children( + &self, + parent_thread_id: ThreadId, + ) -> anyhow::Result> { + self.list_thread_spawn_children_matching(parent_thread_id, /*status*/ None) + .await + } + /// List spawned descendants of `root_thread_id` whose edges match `status`. /// /// Descendants are returned breadth-first by depth, then by thread id for stable ordering. @@ -1871,4 +1880,65 @@ mod tests { .expect("all descendants should load"); assert_eq!(all_descendants, vec![child_thread_id, grandchild_thread_id]); } + + #[tokio::test] + async fn thread_spawn_children_without_status_filter_lists_all_statuses() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should initialize"); + let parent_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000910").expect("valid thread id"); + let open_child_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000911").expect("valid thread id"); + let closed_child_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000912").expect("valid thread id"); + let future_child_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000913").expect("valid thread id"); + + runtime + .upsert_thread_spawn_edge( + parent_thread_id, + open_child_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("open child edge insert should succeed"); + runtime + .upsert_thread_spawn_edge( + parent_thread_id, + closed_child_thread_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + .expect("closed child edge insert should succeed"); + sqlx::query( + r#" +INSERT INTO thread_spawn_edges ( + parent_thread_id, + child_thread_id, + status +) VALUES (?, ?, ?) + "#, + ) + .bind(parent_thread_id.to_string()) + .bind(future_child_thread_id.to_string()) + .bind("future") + .execute(runtime.pool.as_ref()) + .await + .expect("future-status child edge insert should succeed"); + + let children = runtime + .list_thread_spawn_children(parent_thread_id) + .await + .expect("all children should load"); + assert_eq!( + children, + vec![ + open_child_thread_id, + closed_child_thread_id, + future_child_thread_id, + ] + ); + } }