Revert state DB injection and agent graph store (#21481)

## Why

Reverts #20689 to restore the previous optional state DB plumbing. The
conflict resolution keeps the newer installation ID and session/thread
identity changes that landed after #20689, while removing the mandatory
state DB and agent graph store dependency from ThreadManager
construction.

## What changed

- Restored `Option<StateDbHandle>` through app-server, MCP server,
prompt debug, and test entry points.
- Removed the `codex-core` dependency on `codex-agent-graph-store` and
reverted descendant lookup back to the existing state DB path when
available.
- Kept newer `installation_id` forwarding by passing it beside the
optional DB handle.
- Kept local thread-name updates working when the optional state DB
handle is absent.

## Validation

- `git diff --check`
- `cargo test -p codex-thread-store`
- `cargo test -p codex-state -p codex-rollout -p
codex-app-server-protocol`
- Attempted `env CARGO_INCREMENTAL=0 cargo test -p codex-core -p
codex-app-server -p codex-app-server-client -p codex-mcp-server -p
codex-thread-manager-sample -p codex-tui`; blocked locally by a rustc
ICE while compiling `v8 v146.4.0` with `rustc 1.93.0 (254b59607
2026-01-19)` on `aarch64-apple-darwin`.
This commit is contained in:
pakrym-oai
2026-05-06 22:48:29 -07:00
committed by GitHub
Unverified
parent 5bc33fe31f
commit a8488fec5e
54 changed files with 781 additions and 834 deletions
+39 -96
View File
@@ -7,7 +7,6 @@ use crate::environment_selection::default_thread_environment_selections;
use crate::environment_selection::resolve_environment_selections;
use crate::file_watcher::FileWatcher;
use crate::mcp::McpManager;
use crate::resolve_installation_id;
use crate::rollout::RolloutRecorder;
use crate::rollout::truncation;
use crate::session::Codex;
@@ -19,8 +18,6 @@ use crate::skills_watcher::SkillsWatcher;
use crate::skills_watcher::SkillsWatcherEvent;
use crate::tasks::InterruptedTurnHistoryMarker;
use crate::tasks::interrupted_turn_history_marker;
use codex_agent_graph_store::AgentGraphStore;
use codex_agent_graph_store::LocalAgentGraphStore;
use codex_analytics::AnalyticsEventsClient;
use codex_app_server_protocol::ThreadHistoryBuilder;
use codex_app_server_protocol::TurnStatus;
@@ -53,8 +50,8 @@ use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::protocol::W3cTraceContext;
use codex_rollout::state_db;
use codex_rollout::state_db::StateDbHandle;
use codex_state::DirectionalThreadSpawnEdgeStatus;
use codex_thread_store::InMemoryThreadStore;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::LocalThreadStoreConfig;
@@ -251,11 +248,10 @@ pub(crate) struct ThreadManagerState {
mcp_manager: Arc<McpManager>,
skills_watcher: Arc<SkillsWatcher>,
thread_store: Arc<dyn ThreadStore>,
state_db: StateDbHandle,
agent_graph_store: Arc<dyn AgentGraphStore>,
session_source: SessionSource,
installation_id: String,
analytics_events_client: Option<AnalyticsEventsClient>,
state_db: Option<StateDbHandle>,
// Captures submitted ops for testing purpose when test mode is enabled.
ops_log: Option<SharedCapturedOps>,
}
@@ -271,11 +267,10 @@ pub fn build_models_manager(
)
}
pub async fn init_state_db_from_config(config: &Config) -> Option<StateDbHandle> {
state_db::init(config).await
}
pub fn thread_store_from_config(config: &Config, state_db: StateDbHandle) -> Arc<dyn ThreadStore> {
pub fn thread_store_from_config(
config: &Config,
state_db: Option<StateDbHandle>,
) -> Arc<dyn ThreadStore> {
match &config.experimental_thread_store {
ThreadStoreConfig::Local => Arc::new(LocalThreadStore::new(
LocalThreadStoreConfig::from_config(config),
@@ -286,27 +281,6 @@ pub fn thread_store_from_config(config: &Config, state_db: StateDbHandle) -> Arc
}
}
pub fn agent_graph_store_from_state_db(state_db: StateDbHandle) -> Arc<dyn AgentGraphStore> {
Arc::new(LocalAgentGraphStore::new(state_db))
}
async fn state_db_from_roots_for_tests(
codex_home: PathBuf,
sqlite_home: PathBuf,
default_model_provider_id: String,
) -> StateDbHandle {
let config = codex_rollout::RolloutConfig {
codex_home: codex_home.clone(),
sqlite_home,
cwd: codex_home,
model_provider_id: default_model_provider_id,
generate_memories: false,
};
state_db::try_init(&config)
.await
.unwrap_or_else(|err| panic!("test state db should initialize: {err}"))
}
impl ThreadManager {
#[allow(clippy::too_many_arguments)]
pub fn new(
@@ -315,9 +289,8 @@ impl ThreadManager {
session_source: SessionSource,
environment_manager: Arc<EnvironmentManager>,
analytics_events_client: Option<AnalyticsEventsClient>,
state_db: StateDbHandle,
thread_store: Arc<dyn ThreadStore>,
agent_graph_store: Arc<dyn AgentGraphStore>,
state_db: Option<StateDbHandle>,
installation_id: String,
) -> Self {
let codex_home = config.codex_home.clone();
@@ -345,12 +318,11 @@ impl ThreadManager {
mcp_manager,
skills_watcher,
thread_store,
state_db,
agent_graph_store,
auth_manager,
session_source,
installation_id,
analytics_events_client,
state_db,
ops_log: should_use_test_thread_manager_behavior()
.then(|| Arc::new(std::sync::Mutex::new(Vec::new()))),
}),
@@ -360,7 +332,7 @@ impl ThreadManager {
/// Construct with a dummy AuthManager containing the provided CodexAuth.
/// Used for integration tests: should not be used by ordinary business logic.
pub(crate) async fn with_models_provider_for_tests(
pub(crate) fn with_models_provider_for_tests(
auth: CodexAuth,
provider: ModelProviderInfo,
) -> Self {
@@ -371,27 +343,11 @@ impl ThreadManager {
));
std::fs::create_dir_all(&codex_home)
.unwrap_or_else(|err| panic!("temp codex home dir create failed: {err}"));
let state_db = state_db_from_roots_for_tests(
codex_home.clone(),
codex_home.clone(),
OPENAI_PROVIDER_ID.to_string(),
)
.await;
let skills_codex_home = match AbsolutePathBuf::from_absolute_path_checked(&codex_home) {
Ok(codex_home) => codex_home,
Err(err) => panic!("test codex_home should be absolute: {err}"),
};
let installation_id = resolve_installation_id(&skills_codex_home)
.await
.unwrap_or_else(|err| panic!("resolve test installation id failed: {err}"));
let mut manager = Self::with_models_provider_and_home_and_state_db_for_tests(
let mut manager = Self::with_models_provider_and_home_for_tests(
auth,
provider,
codex_home.clone(),
Arc::new(EnvironmentManager::default_for_tests()),
state_db,
skills_codex_home,
installation_id,
);
manager._test_codex_home_guard = Some(TempCodexHomeGuard { path: codex_home });
manager
@@ -399,47 +355,35 @@ impl ThreadManager {
/// Construct with a dummy AuthManager containing the provided CodexAuth and codex home.
/// Used for integration tests: should not be used by ordinary business logic.
pub(crate) async fn with_models_provider_and_home_for_tests(
pub(crate) fn with_models_provider_and_home_for_tests(
auth: CodexAuth,
provider: ModelProviderInfo,
codex_home: PathBuf,
environment_manager: Arc<EnvironmentManager>,
) -> Self {
let state_db = state_db_from_roots_for_tests(
codex_home.clone(),
codex_home.clone(),
OPENAI_PROVIDER_ID.to_string(),
)
.await;
let skills_codex_home = match AbsolutePathBuf::from_absolute_path_checked(&codex_home) {
Ok(codex_home) => codex_home,
Err(err) => panic!("test codex_home should be absolute: {err}"),
};
let installation_id = resolve_installation_id(&skills_codex_home)
.await
.unwrap_or_else(|err| panic!("resolve test installation id failed: {err}"));
Self::with_models_provider_and_home_and_state_db_for_tests(
Self::with_models_provider_home_and_state_for_tests(
auth,
provider,
codex_home,
environment_manager,
state_db,
skills_codex_home,
installation_id,
/*state_db*/ None,
)
}
fn with_models_provider_and_home_and_state_db_for_tests(
pub(crate) fn with_models_provider_home_and_state_for_tests(
auth: CodexAuth,
provider: ModelProviderInfo,
codex_home: PathBuf,
environment_manager: Arc<EnvironmentManager>,
state_db: StateDbHandle,
skills_codex_home: AbsolutePathBuf,
installation_id: String,
state_db: Option<StateDbHandle>,
) -> Self {
set_thread_manager_test_mode_for_tests(/*enabled*/ true);
let auth_manager = AuthManager::from_auth_for_testing(auth);
let installation_id = uuid::Uuid::new_v4().to_string();
let skills_codex_home = match AbsolutePathBuf::from_absolute_path_checked(&codex_home) {
Ok(codex_home) => codex_home,
Err(err) => panic!("test codex_home should be absolute: {err}"),
};
let (thread_created_tx, _) = broadcast::channel(THREAD_CREATED_CHANNEL_CAPACITY);
let restriction_product = SessionSource::Exec.restriction_product();
let plugins_manager = Arc::new(PluginsManager::new_with_restriction_product(
@@ -458,11 +402,11 @@ impl ThreadManager {
let thread_store: Arc<dyn ThreadStore> = Arc::new(LocalThreadStore::new(
LocalThreadStoreConfig {
codex_home: codex_home.clone(),
sqlite_home: codex_home.clone(),
default_model_provider_id: OPENAI_PROVIDER_ID.to_string(),
},
state_db.clone(),
));
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
Self {
state: Arc::new(ThreadManagerState {
threads: Arc::new(RwLock::new(HashMap::new())),
@@ -475,12 +419,11 @@ impl ThreadManager {
mcp_manager,
skills_watcher,
thread_store,
state_db,
agent_graph_store,
auth_manager,
session_source: SessionSource::Exec,
installation_id,
analytics_events_client: None,
state_db,
ops_log: should_use_test_thread_manager_behavior()
.then(|| Arc::new(std::sync::Mutex::new(Vec::new()))),
}),
@@ -566,17 +509,22 @@ impl ThreadManager {
subtree_thread_ids.push(thread_id);
seen_thread_ids.insert(thread_id);
for descendant_id in self
.state
.agent_graph_store
.list_thread_spawn_descendants(thread_id, /*status_filter*/ None)
.await
.map_err(|err| {
CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}"))
})?
{
if seen_thread_ids.insert(descendant_id) {
subtree_thread_ids.push(descendant_id);
if let Some(state_db_ctx) = thread.state_db() {
for status in [
DirectionalThreadSpawnEdgeStatus::Open,
DirectionalThreadSpawnEdgeStatus::Closed,
] {
for descendant_id in state_db_ctx
.list_thread_spawn_descendants_with_status(thread_id, status)
.await
.map_err(|err| {
CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}"))
})?
{
if seen_thread_ids.insert(descendant_id) {
subtree_thread_ids.push(descendant_id);
}
}
}
}
@@ -918,14 +866,10 @@ impl ThreadManager {
}
impl ThreadManagerState {
pub(crate) fn state_db(&self) -> StateDbHandle {
pub(crate) fn state_db(&self) -> Option<StateDbHandle> {
self.state_db.clone()
}
pub(crate) fn agent_graph_store(&self) -> Arc<dyn AgentGraphStore> {
self.agent_graph_store.clone()
}
pub(crate) async fn list_thread_ids(&self) -> Vec<ThreadId> {
self.threads
.read()
@@ -1242,7 +1186,6 @@ impl ThreadManagerState {
parent_trace,
environment_selections,
analytics_events_client: self.analytics_events_client.clone(),
state_db: Some(self.state_db.clone()),
thread_store: Arc::clone(&self.thread_store),
})
.await?;