agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
Tom
2026-06-24 13:24:10 -07:00
committed by GitHub
Unverified
parent 989f55defa
commit ece1dfece0
26 changed files with 317 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/thread_manager_tests.rs
+95 -11
View File
@@ -35,6 +35,49 @@ use wiremock::MockServer;
const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111";
struct FakeAgentGraphStore {
root_thread_id: ThreadId,
descendant_thread_ids: Vec<ThreadId>,
}
impl codex_agent_graph_store::AgentGraphStore for FakeAgentGraphStore {
fn upsert_thread_spawn_edge(
&self,
_parent_thread_id: ThreadId,
_child_thread_id: ThreadId,
_status: codex_agent_graph_store::ThreadSpawnEdgeStatus,
) -> codex_agent_graph_store::AgentGraphStoreFuture<'_, ()> {
Box::pin(async { panic!("unexpected graph upsert") })
}
fn set_thread_spawn_edge_status(
&self,
_child_thread_id: ThreadId,
_status: codex_agent_graph_store::ThreadSpawnEdgeStatus,
) -> codex_agent_graph_store::AgentGraphStoreFuture<'_, ()> {
Box::pin(async { panic!("unexpected graph status update") })
}
fn list_thread_spawn_children(
&self,
_parent_thread_id: ThreadId,
_status_filter: Option<codex_agent_graph_store::ThreadSpawnEdgeStatus>,
) -> codex_agent_graph_store::AgentGraphStoreFuture<'_, Vec<ThreadId>> {
Box::pin(async { panic!("unexpected direct-child listing") })
}
fn list_thread_spawn_descendants(
&self,
root_thread_id: ThreadId,
status_filter: Option<codex_agent_graph_store::ThreadSpawnEdgeStatus>,
) -> codex_agent_graph_store::AgentGraphStoreFuture<'_, Vec<ThreadId>> {
assert_eq!(root_thread_id, self.root_thread_id);
assert_eq!(status_filter, None);
let descendant_thread_ids = self.descendant_thread_ids.clone();
Box::pin(async move { Ok(descendant_thread_ids) })
}
}
fn user_msg(text: &str) -> ResponseItem {
ResponseItem::Message {
id: None,
@@ -477,7 +520,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
/*agent_graph_store*/ None,
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -588,7 +631,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() {
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
/*agent_graph_store*/ None,
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -713,7 +756,7 @@ async fn explicit_installation_id_skips_codex_home_file() {
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store,
state_db.clone(),
local_agent_graph_store_from_state_db(state_db.as_ref()),
installation_id.clone(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -754,7 +797,7 @@ async fn resume_active_thread_from_rollout_returns_running_thread() {
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
/*agent_graph_store*/ None,
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -814,7 +857,7 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() {
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
/*agent_graph_store*/ None,
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -881,7 +924,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() {
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store,
state_db.clone(),
local_agent_graph_store_from_state_db(state_db.as_ref()),
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -948,6 +991,47 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() {
.expect("shutdown resumed thread");
}
#[tokio::test]
async fn subtree_listing_uses_injected_graph_store_without_state_db() {
let temp_dir = tempdir().expect("tempdir");
let mut config = test_config().await;
config.codex_home = temp_dir.path().join("codex-home").abs();
std::fs::create_dir_all(&config.codex_home).expect("create codex home");
let root_thread_id = ThreadId::new();
let descendant_thread_ids = vec![ThreadId::new(), ThreadId::new()];
let agent_graph_store = Arc::new(FakeAgentGraphStore {
root_thread_id,
descendant_thread_ids: descendant_thread_ids.clone(),
});
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let manager = ThreadManager::new(
&config,
auth_manager,
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
Some(agent_graph_store),
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
);
let mut expected_thread_ids = vec![root_thread_id];
expected_thread_ids.extend(descendant_thread_ids);
assert_eq!(
manager
.list_agent_subtree_thread_ids(root_thread_id)
.await
.expect("subtree should load from injected graph store"),
expected_thread_ids
);
}
#[tokio::test]
async fn rollout_path_resume_and_fork_read_history_through_thread_store() {
let temp_dir = tempdir().expect("tempdir");
@@ -976,7 +1060,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() {
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store.clone(),
state_db,
local_agent_graph_store_from_state_db(state_db.as_ref()),
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -1081,7 +1165,7 @@ async fn new_uses_active_provider_for_model_refresh() {
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
/*agent_graph_store*/ None,
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -1303,7 +1387,7 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
local_agent_graph_store_from_state_db(state_db.as_ref()),
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -1412,7 +1496,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() {
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
local_agent_graph_store_from_state_db(state_db.as_ref()),
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
@@ -1511,7 +1595,7 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
local_agent_graph_store_from_state_db(state_db.as_ref()),
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,