Files
codex/codex-rs/core/src/personality_migration.rs
T
Rasmus Rygaard 7e310bc7f3 Inject state DB, agent graph store (#20689)
## Why

We want the agent graph store to be passed down the stack as a real
dependency, the same way we already treat the thread store.

This will let us inject the agent graph store as a real dependency and
support implementations other than the local SQLite-backed one. Right
now most code instantiates a state DB and an agent graph store
just-in-time. Ideally, we would not depend on the state DB directly but
only read through the higher-level interfaces.

This change makes the dependency boundaries explicit and moves state DB
initialization to process bootstrap instead of hiding it inside local
store implementations.

## What changed

- `ThreadManager` now requires a `StateDbHandle` and an
`AgentGraphStore` at construction time instead of treating them as
optional internals.
- The local store constructors no longer lazily initialize SQLite.
Callers now initialize the state DB once per process and use that shared
handle to build:
  - `LocalThreadStore`
  - `LocalAgentGraphStore`
- App bootstraps (`app-server`, `mcp-server`, `prompt_debug`, and the
thread-manager sample) now initialize the state DB up front and inject
the resulting handle down the stack.
- `app-server` now consistently uses its process-scoped state DB handle
instead of reopening SQLite or trying to recover it from loaded threads.
- Device-key storage now reuses the shared state DB handle instead of
maintaining its own lazy opener.
- The thread archive / descendant traversal paths now use the injected
`AgentGraphStore` instead of reaching through local
thread-store-specific state.

## Verification

- `cargo check -p codex-core -p codex-thread-store -p codex-app-server
-p codex-mcp-server -p codex-thread-manager-sample --tests`
- `cargo test -p codex-thread-store`
- `cargo test -p codex-core
thread_manager_accepts_separate_agent_graph_store_and_thread_store --
--nocapture`
- `cargo test -p codex-app-server
thread_archive_archives_spawned_descendants -- --nocapture`
2026-05-05 21:45:29 +00:00

116 lines
3.6 KiB
Rust

use crate::config::edit::ConfigEditsBuilder;
use codex_config::config_toml::ConfigToml;
use codex_protocol::config_types::Personality;
use codex_rollout::state_db::StateDbHandle;
use codex_thread_store::ListThreadsParams;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::LocalThreadStoreConfig;
use codex_thread_store::ThreadSortKey;
use codex_thread_store::ThreadStore;
use std::io;
use std::path::Path;
use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;
pub const PERSONALITY_MIGRATION_FILENAME: &str = ".personality_migration";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PersonalityMigrationStatus {
SkippedMarker,
SkippedExplicitPersonality,
SkippedNoSessions,
Applied,
}
pub async fn maybe_migrate_personality(
codex_home: &Path,
config_toml: &ConfigToml,
state_db: StateDbHandle,
) -> io::Result<PersonalityMigrationStatus> {
let marker_path = codex_home.join(PERSONALITY_MIGRATION_FILENAME);
if tokio::fs::try_exists(&marker_path).await? {
return Ok(PersonalityMigrationStatus::SkippedMarker);
}
let config_profile = config_toml
.get_config_profile(/*override_profile*/ None)
.map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
if config_toml.personality.is_some() || config_profile.personality.is_some() {
create_marker(&marker_path).await?;
return Ok(PersonalityMigrationStatus::SkippedExplicitPersonality);
}
let model_provider_id = config_profile
.model_provider
.or_else(|| config_toml.model_provider.clone())
.unwrap_or_else(|| "openai".to_string());
if !has_recorded_sessions(codex_home, model_provider_id.as_str(), state_db).await? {
create_marker(&marker_path).await?;
return Ok(PersonalityMigrationStatus::SkippedNoSessions);
}
ConfigEditsBuilder::new(codex_home)
.set_personality(Some(Personality::Pragmatic))
.apply()
.await
.map_err(|err| {
io::Error::other(format!("failed to persist personality migration: {err}"))
})?;
create_marker(&marker_path).await?;
Ok(PersonalityMigrationStatus::Applied)
}
async fn has_recorded_sessions(
codex_home: &Path,
default_provider: &str,
state_db: StateDbHandle,
) -> io::Result<bool> {
let config = LocalThreadStoreConfig {
codex_home: codex_home.to_path_buf(),
default_model_provider_id: default_provider.to_string(),
};
let store = LocalThreadStore::new(config, state_db);
if has_threads(&store, /*archived*/ false).await? {
return Ok(true);
}
has_threads(&store, /*archived*/ true).await
}
async fn has_threads(store: &LocalThreadStore, archived: bool) -> io::Result<bool> {
store
.list_threads(ListThreadsParams {
page_size: 1,
cursor: None,
sort_key: ThreadSortKey::CreatedAt,
sort_direction: codex_thread_store::SortDirection::Desc,
allowed_sources: Vec::new(),
model_providers: None,
cwd_filters: None,
archived,
search_term: None,
use_state_db_only: false,
})
.await
.map(|page| !page.items.is_empty())
.map_err(io::Error::other)
}
async fn create_marker(marker_path: &Path) -> io::Result<()> {
match OpenOptions::new()
.create_new(true)
.write(true)
.open(marker_path)
.await
{
Ok(mut file) => file.write_all(b"v1\n").await,
Err(err) if err.kind() == io::ErrorKind::AlreadyExists => Ok(()),
Err(err) => Err(err),
}
}
#[cfg(test)]
#[path = "personality_migration_tests.rs"]
mod tests;