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

Inject state DB, agent graph store (#20689)

Browse Source 
## 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`
This commit is contained in:
Rasmus Rygaard
2026-05-05 21:45:29 +00:00
committed by GitHub
parent 36460387ec
commit 7e310bc7f3
53 changed files with 768 additions and 745 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/thread-store/src/local/archive_thread.rs
+11 -19
View File
@@ -13,11 +13,11 @@ pub(super) async fn archive_thread(
params: ArchiveThreadParams,
) -> ThreadStoreResult<()> {
let thread_id = params.thread_id;
let state_db_ctx = store.state_db().await;
let state_db = store.state_db();
let rollout_path = find_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -52,11 +52,10 @@ pub(super) async fn archive_thread(
}
})?;
if let Some(ctx) = state_db_ctx {
let _ = ctx
.mark_archived(thread_id, archived_path.as_path(), Utc::now())
.await;
}
let _ = store
.state_db()
.mark_archived(thread_id, archived_path.as_path(), Utc::now())
.await;
Ok(())
}
@@ -75,13 +74,15 @@ mod tests {
use crate::ThreadSortKey;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_session_file;
#[tokio::test]
async fn archive_thread_moves_rollout_to_archived_collection() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(201);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
@@ -127,21 +128,12 @@ mod tests {
async fn archive_thread_updates_sqlite_metadata_when_present() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(202);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
let mut builder = codex_state::ThreadMetadataBuilder::new(
thread_id,
active_path.clone(),
codex-rs/thread-store/src/local/create_thread.rs
+2 -2
View File
@@ -22,12 +22,12 @@ pub(super) async fn create_thread(
})?;
let config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
sqlite_home: store.sqlite_home(),
cwd,
model_provider_id: params.metadata.model_provider.clone(),
generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled),
};
let state_db_ctx = store.state_db().await;
let state_db_ctx = Some(store.state_db());
let recorder = RolloutRecorder::new(
&config,
RolloutRecorderParams::new(
codex-rs/thread-store/src/local/list_threads.rs
+32 -36
View File
@@ -39,16 +39,16 @@ pub(super) async fn list_threads(
SortDirection::Asc => codex_rollout::SortDirection::Asc,
SortDirection::Desc => codex_rollout::SortDirection::Desc,
};
let state_db = store.state_db().await;
let rollout_config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
sqlite_home: store.sqlite_home(),
cwd: store.config.codex_home.clone(),
model_provider_id: store.config.default_model_provider_id.clone(),
generate_memories: false,
};
let state_db_ctx = Some(store.state_db());
let page = list_rollout_threads(
state_db,
state_db_ctx,
&rollout_config,
store.config.default_model_provider_id.as_str(),
&params,
@@ -80,14 +80,13 @@ pub(super) async fn list_threads(
.map(|thread| thread.thread_id)
.collect::<HashSet<_>>();
let mut names = HashMap::<ThreadId, String>::with_capacity(thread_ids.len());
if let Some(state_db_ctx) = store.state_db().await {
for &thread_id in &thread_ids {
let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else {
continue;
};
if let Some(title) = distinct_thread_metadata_title(&metadata) {
names.insert(thread_id, title);
}
let state_db_ctx = store.state_db();
for &thread_id in &thread_ids {
let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else {
continue;
};
if let Some(title) = distinct_thread_metadata_title(&metadata) {
names.insert(thread_id, title);
}
}
if names.len() < thread_ids.len()
@@ -108,9 +107,9 @@ pub(super) async fn list_threads(
}
async fn list_rollout_threads(
state_db: Option<codex_rollout::StateDbHandle>,
state_db_ctx: Option<codex_rollout::StateDbHandle>,
config: &RolloutConfig,
default_model_provider_id: &str,
default_model_provider: &str,
params: &ListThreadsParams,
cursor: Option<&codex_rollout::Cursor>,
sort_key: codex_rollout::ThreadSortKey,
@@ -118,7 +117,7 @@ async fn list_rollout_threads(
) -> ThreadStoreResult<codex_rollout::ThreadsPage> {
let page = if params.use_state_db_only && params.archived {
RolloutRecorder::list_archived_threads_from_state_db(
state_db,
state_db_ctx.clone(),
config,
params.page_size,
cursor,
@@ -127,13 +126,13 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
default_model_provider_id,
default_model_provider,
params.search_term.as_deref(),
)
.await
} else if params.use_state_db_only {
RolloutRecorder::list_threads_from_state_db(
state_db,
state_db_ctx.clone(),
config,
params.page_size,
cursor,
@@ -142,13 +141,13 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
default_model_provider_id,
default_model_provider,
params.search_term.as_deref(),
)
.await
} else if params.archived {
RolloutRecorder::list_archived_threads(
state_db,
state_db_ctx.clone(),
config,
params.page_size,
cursor,
@@ -157,13 +156,13 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
default_model_provider_id,
default_model_provider,
params.search_term.as_deref(),
)
.await
} else {
RolloutRecorder::list_threads(
state_db,
state_db_ctx,
config,
params.page_size,
cursor,
@@ -172,7 +171,7 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
default_model_provider_id,
default_model_provider,
params.search_term.as_deref(),
)
.await
@@ -195,7 +194,9 @@ mod tests {
use super::*;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
use crate::local::test_support::write_session_file_with;
@@ -203,7 +204,7 @@ mod tests {
#[tokio::test]
async fn list_threads_uses_default_provider_when_rollout_omits_provider() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
write_session_file_with(
home.path(),
home.path().join("sessions/2025/01/03"),
@@ -238,22 +239,13 @@ mod tests {
async fn list_threads_preserves_sqlite_title_search_results() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(103);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = home.path().join("rollout-title-search.jsonl");
fs::write(&rollout_path, "").expect("placeholder rollout file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
let created_at = Utc::now();
let mut builder = codex_state::ThreadMetadataBuilder::new(
thread_id,
@@ -267,6 +259,10 @@ mod tests {
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.title = "needle title".to_string();
metadata.first_user_message = Some("plain preview".to_string());
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
runtime
.upsert_thread(&metadata)
.await
@@ -303,7 +299,7 @@ mod tests {
#[tokio::test]
async fn list_threads_selects_active_or_archived_collection() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let active_uuid = Uuid::from_u128(105);
let archived_uuid = Uuid::from_u128(106);
write_session_file(home.path(), "2025-01-03T12-00-00", active_uuid)
@@ -372,7 +368,7 @@ mod tests {
async fn list_threads_returns_local_rollout_summary() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let store = LocalThreadStore::new(config, /*state_db*/ None);
let store = LocalThreadStore::new(config.clone(), init_test_state_db(&config).await);
let uuid = Uuid::from_u128(101);
let path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
@@ -411,7 +407,7 @@ mod tests {
#[tokio::test]
async fn list_threads_rejects_invalid_cursor() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let err = store
.list_threads(ListThreadsParams {
codex-rs/thread-store/src/local/live_writer.rs
+2 -2
View File
@@ -66,12 +66,12 @@ pub(super) async fn resume_thread(
})?;
let config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
sqlite_home: store.sqlite_home(),
cwd,
model_provider_id: params.metadata.model_provider.clone(),
generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled),
};
let state_db_ctx = store.state_db().await;
let state_db_ctx = Some(store.state_db());
let recorder = RolloutRecorder::new(
&config,
RolloutRecorderParams::resume(
codex-rs/thread-store/src/local/mod.rs
+24 -18
View File
@@ -41,7 +41,7 @@ use crate::UpdateThreadMetadataParams;
pub struct LocalThreadStore {
pub(super) config: LocalThreadStoreConfig,
live_recorders: Arc<Mutex<HashMap<ThreadId, RolloutRecorder>>>,
state_db: Option<StateDbHandle>,
state_db: StateDbHandle,
}
/// Process-scoped configuration for local thread storage.
@@ -51,7 +51,6 @@ pub struct LocalThreadStore {
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LocalThreadStoreConfig {
pub codex_home: PathBuf,
pub sqlite_home: PathBuf,
/// Provider used only when older local metadata does not contain one.
pub default_model_provider_id: String,
}
@@ -60,7 +59,6 @@ impl LocalThreadStoreConfig {
pub fn from_config(config: &impl codex_rollout::RolloutConfigView) -> Self {
Self {
codex_home: config.codex_home().to_path_buf(),
sqlite_home: config.sqlite_home().to_path_buf(),
default_model_provider_id: config.model_provider_id().to_string(),
}
}
@@ -75,8 +73,9 @@ impl std::fmt::Debug for LocalThreadStore {
}
impl LocalThreadStore {
/// Create a local store using an already initialized state DB handle.
pub fn new(config: LocalThreadStoreConfig, state_db: Option<StateDbHandle>) -> Self {
/// Create a local store from process-scoped local storage configuration and
/// the caller-provided shared state DB handle.
pub fn new(config: LocalThreadStoreConfig, state_db: StateDbHandle) -> Self {
Self {
config,
live_recorders: Arc::new(Mutex::new(HashMap::new())),
@@ -85,10 +84,14 @@ impl LocalThreadStore {
}
/// Return the state DB handle used by local rollout writers.
pub async fn state_db(&self) -> Option<StateDbHandle> {
pub fn state_db(&self) -> StateDbHandle {
self.state_db.clone()
}
pub(super) fn sqlite_home(&self) -> PathBuf {
self.state_db.codex_home().to_path_buf()
}
/// Read a local rollout-backed thread by path.
pub async fn read_thread_by_rollout_path(
&self,
@@ -282,14 +285,16 @@ mod tests {
use super::*;
use crate::ThreadEventPersistenceMode;
use crate::ThreadPersistenceMetadata;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
#[tokio::test]
async fn live_writer_lifecycle_writes_and_closes() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
@@ -338,7 +343,7 @@ mod tests {
#[tokio::test]
async fn create_thread_rejects_missing_cwd() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
let mut params = create_thread_params(thread_id);
params.metadata.cwd = None;
@@ -358,7 +363,7 @@ mod tests {
#[tokio::test]
async fn discard_thread_drops_unmaterialized_live_writer() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
@@ -396,8 +401,9 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let thread_id = ThreadId::default();
let state_db = init_test_state_db(&config).await;
let first_store = LocalThreadStore::new(config.clone(), /*state_db*/ None);
let first_store = LocalThreadStore::new(config.clone(), state_db.clone());
first_store
.create_thread(create_thread_params(thread_id))
.await
@@ -426,7 +432,7 @@ mod tests {
.await
.expect("shutdown initial writer");
let resumed_store = LocalThreadStore::new(config, /*state_db*/ None);
let resumed_store = LocalThreadStore::new(config, state_db);
resumed_store
.resume_thread(ResumeThreadParams {
thread_id,
@@ -457,7 +463,7 @@ mod tests {
#[tokio::test]
async fn create_thread_rejects_duplicate_live_writer() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
@@ -477,7 +483,7 @@ mod tests {
#[tokio::test]
async fn resume_thread_rejects_duplicate_live_writer() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
@@ -506,7 +512,7 @@ mod tests {
#[tokio::test]
async fn resume_thread_rejects_missing_cwd() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = uuid::Uuid::from_u128(407);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
@@ -535,7 +541,7 @@ mod tests {
async fn load_history_uses_live_writer_rollout_path() {
let home = TempDir::new().expect("temp dir");
let external_home = TempDir::new().expect("external temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = uuid::Uuid::from_u128(404);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = write_session_file(external_home.path(), "2025-01-04T10-00-00", uuid)
@@ -584,7 +590,7 @@ mod tests {
async fn read_thread_uses_live_writer_rollout_path_for_external_resume() {
let home = TempDir::new().expect("temp dir");
let external_home = TempDir::new().expect("external temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = uuid::Uuid::from_u128(406);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = write_session_file(external_home.path(), "2025-01-04T11-00-00", uuid)
@@ -623,7 +629,7 @@ mod tests {
#[tokio::test]
async fn load_history_uses_live_writer_rollout_path_for_archived_source() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = uuid::Uuid::from_u128(405);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = write_archived_session_file(home.path(), "2025-01-04T10-30-00", uuid)
@@ -691,7 +697,7 @@ mod tests {
#[tokio::test]
async fn read_thread_by_rollout_path_includes_history() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
codex-rs/thread-store/src/local/read_thread.rs
+33 -77
View File
@@ -176,12 +176,12 @@ async fn resolve_rollout_path(
return Ok(Some(path));
}
let state_db_ctx = store.state_db().await;
let state_db = store.state_db();
if include_archived {
match find_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -191,7 +191,7 @@ async fn resolve_rollout_path(
None => find_archived_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -202,7 +202,7 @@ async fn resolve_rollout_path(
find_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -260,8 +260,7 @@ async fn read_sqlite_metadata(
store: &LocalThreadStore,
thread_id: codex_protocol::ThreadId,
) -> Option<ThreadMetadata> {
let runtime = store.state_db().await?;
runtime.get_thread(thread_id).await.ok().flatten()
store.state_db().get_thread(thread_id).await.ok().flatten()
}
async fn stored_thread_from_sqlite_metadata(
@@ -412,7 +411,9 @@ mod tests {
use super::*;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
use crate::local::test_support::write_session_file_with_fork;
@@ -420,7 +421,7 @@ mod tests {
#[tokio::test]
async fn read_thread_returns_active_rollout_summary() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(205);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
@@ -448,7 +449,7 @@ mod tests {
#[tokio::test]
async fn read_thread_returns_rollout_path_summary() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(211);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
@@ -479,17 +480,12 @@ mod tests {
async fn read_thread_by_rollout_path_prefers_sqlite_git_info() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(223);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
active_path.clone(),
@@ -527,7 +523,7 @@ mod tests {
#[tokio::test]
async fn read_thread_returns_archived_rollout_when_requested() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(207);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid)
@@ -568,7 +564,7 @@ mod tests {
#[tokio::test]
async fn read_thread_prefers_active_rollout_over_archived() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(208);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
@@ -593,7 +589,7 @@ mod tests {
#[tokio::test]
async fn read_thread_returns_forked_from_id() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(209);
let parent_uuid = Uuid::from_u128(210);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
@@ -626,17 +622,12 @@ mod tests {
async fn read_thread_applies_sqlite_thread_name() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(212);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder =
ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli);
builder.model_provider = Some(config.default_model_provider_id.clone());
@@ -666,13 +657,8 @@ mod tests {
async fn read_thread_preserves_rollout_cwd_when_sqlite_metadata_exists() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(224);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let day_dir = home.path().join("sessions/2025/01/03");
@@ -741,7 +727,7 @@ mod tests {
#[tokio::test]
async fn read_thread_uses_legacy_thread_name_when_sqlite_title_is_missing() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(213);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
@@ -765,6 +751,8 @@ mod tests {
async fn read_thread_uses_sqlite_metadata_for_rollout_without_user_preview() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(217);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let day_dir = home.path().join("sessions/2025/01/03");
@@ -786,13 +774,6 @@ mod tests {
});
writeln!(file, "{meta}").expect("write session meta");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
rollout_path.clone(),
@@ -835,18 +816,13 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let external = TempDir::new().expect("external temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(220);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
let stale_path = external.path().join("missing-rollout.jsonl");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
stale_path.clone(),
@@ -884,6 +860,8 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let external = TempDir::new().expect("external temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(221);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
@@ -891,13 +869,6 @@ mod tests {
let other_uuid = Uuid::from_u128(222);
let stale_path = write_session_file(external.path(), "2025-01-04T12-00-00", other_uuid)
.expect("other session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder =
ThreadMetadataBuilder::new(thread_id, stale_path, Utc::now(), SessionSource::Cli);
builder.model_provider = Some("wrong-sqlite-provider".to_string());
@@ -929,7 +900,7 @@ mod tests {
#[tokio::test]
async fn read_thread_uses_session_meta_for_rollout_without_user_preview_or_sqlite_metadata() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(218);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let day_dir = home.path().join("sessions/2025/01/03");
@@ -984,18 +955,13 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let external = TempDir::new().expect("external temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(214);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = external
.path()
.join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl"));
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
rollout_path.clone(),
@@ -1042,20 +1008,15 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let external = TempDir::new().expect("external temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(216);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = external
.path()
.join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl"));
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let mut builder =
ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli);
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
builder.archived_at = Some(Utc::now());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.first_user_message = Some("Archived SQLite preview".to_string());
@@ -1098,17 +1059,12 @@ mod tests {
async fn read_thread_sqlite_fallback_loads_archived_history() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(219);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
archived_path.clone(),
@@ -1144,7 +1100,7 @@ mod tests {
#[tokio::test]
async fn read_thread_fails_without_rollout() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(206);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
codex-rs/thread-store/src/local/test_support.rs
+17 -1
View File
@@ -4,18 +4,34 @@ use std::path::Path;
use std::path::PathBuf;
use codex_rollout::ARCHIVED_SESSIONS_SUBDIR;
use codex_rollout::StateDbHandle;
use uuid::Uuid;
use super::LocalThreadStore;
use super::LocalThreadStoreConfig;
pub(super) fn test_config(codex_home: &Path) -> LocalThreadStoreConfig {
LocalThreadStoreConfig {
codex_home: codex_home.to_path_buf(),
sqlite_home: codex_home.to_path_buf(),
default_model_provider_id: "test-provider".to_string(),
}
}
pub(super) async fn init_test_state_db(config: &LocalThreadStoreConfig) -> StateDbHandle {
codex_state::StateRuntime::init(
config.codex_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize")
}
pub(super) async fn test_store(codex_home: &Path) -> LocalThreadStore {
let config = test_config(codex_home);
let state_db = init_test_state_db(&config).await;
LocalThreadStore::new(config, state_db)
}
pub(super) fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result<PathBuf> {
write_session_file_with(
root,
codex-rs/thread-store/src/local/unarchive_thread.rs
+11 -19
View File
@@ -17,11 +17,11 @@ pub(super) async fn unarchive_thread(
params: ArchiveThreadParams,
) -> ThreadStoreResult<StoredThread> {
let thread_id = params.thread_id;
let state_db_ctx = store.state_db().await;
let state_db = store.state_db();
let archived_path = find_archived_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -73,11 +73,10 @@ pub(super) async fn unarchive_thread(
message: format!("failed to update unarchived thread timestamp: {err}"),
})?;
if let Some(ctx) = state_db_ctx {
let _ = ctx
.mark_unarchived(thread_id, restored_path.as_path())
.await;
}
let _ = store
.state_db()
.mark_unarchived(thread_id, restored_path.as_path())
.await;
let item = read_thread_item_from_rollout(restored_path.clone())
.await
@@ -112,13 +111,15 @@ mod tests {
use super::*;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
#[tokio::test]
async fn unarchive_thread_restores_rollout_and_returns_updated_thread() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(203);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid)
@@ -149,21 +150,12 @@ mod tests {
async fn unarchive_thread_updates_sqlite_metadata_when_present() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(204);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
let mut builder = codex_state::ThreadMetadataBuilder::new(
thread_id,
archived_path.clone(),
codex-rs/thread-store/src/local/update_thread_metadata.rs
+26 -77
View File
@@ -59,9 +59,8 @@ pub(super) async fn update_thread_metadata(
.await?;
}
let state_db_ctx = store.state_db().await;
codex_rollout::state_db::reconcile_rollout(
state_db_ctx.as_deref(),
Some(store.state_db()).as_deref(),
resolved_rollout_path.path.as_path(),
store.config.default_model_provider_id.as_str(),
/*builder*/ None,
@@ -73,11 +72,7 @@ pub(super) async fn update_thread_metadata(
let resolved_git_info = match git_info {
Some(git_info) => {
let Some(state_db) = store.state_db().await else {
return Err(ThreadStoreError::Internal {
message: format!("sqlite state db unavailable for thread {thread_id}"),
});
};
let state_db = store.state_db();
let metadata =
state_db
.get_thread(thread_id)
@@ -157,11 +152,7 @@ async fn apply_thread_git_info(
branch: &Option<String>,
origin_url: &Option<String>,
) -> ThreadStoreResult<()> {
let Some(state_db) = store.state_db().await else {
return Err(ThreadStoreError::Internal {
message: format!("sqlite state db unavailable for thread {thread_id}"),
});
};
let state_db = store.state_db();
let updated = state_db
.update_thread_git_info(
thread_id,
@@ -307,11 +298,11 @@ async fn resolve_rollout_path(
return Ok(ResolvedRolloutPath { path, archived });
}
let state_db_ctx = store.state_db().await;
let state_db = store.state_db();
let active_path = find_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -331,7 +322,7 @@ async fn resolve_rollout_path(
find_archived_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -366,14 +357,16 @@ mod tests {
use crate::ThreadPersistenceMetadata;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
#[tokio::test]
async fn update_thread_metadata_sets_name_on_active_rollout_and_indexes_name() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(301);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
@@ -408,18 +401,12 @@ mod tests {
async fn update_thread_metadata_sets_memory_mode_on_active_rollout() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(302);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
write_session_file(home.path(), "2025-01-03T14-30-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let thread = store
.update_thread_metadata(UpdateThreadMetadataParams {
thread_id,
@@ -452,13 +439,8 @@ mod tests {
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
write_session_file(home.path(), "2025-01-03T18-30-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
store
.update_thread_metadata(UpdateThreadMetadataParams {
@@ -517,7 +499,7 @@ mod tests {
async fn update_thread_metadata_uses_live_rollout_path_for_external_resume() {
let home = TempDir::new().expect("temp dir");
let external_home = TempDir::new().expect("external temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(307);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path = write_session_file(external_home.path(), "2025-01-03T14-45-00", uuid)
@@ -558,13 +540,8 @@ mod tests {
async fn update_thread_metadata_sets_git_info() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config, Some(runtime));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config, runtime);
let uuid = Uuid::from_u128(309);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
write_session_file(home.path(), "2025-01-03T17-00-00", uuid).expect("session file");
@@ -601,13 +578,8 @@ mod tests {
async fn update_thread_metadata_partially_updates_git_info() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config, Some(runtime));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config, runtime);
let uuid = Uuid::from_u128(310);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
write_session_file(home.path(), "2025-01-03T17-30-00", uuid).expect("session file");
@@ -659,13 +631,8 @@ mod tests {
async fn update_thread_metadata_clears_git_info_fields() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(311);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
@@ -829,7 +796,7 @@ mod tests {
#[tokio::test]
async fn update_thread_metadata_rejects_mismatched_session_meta_id() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let filename_uuid = Uuid::from_u128(303);
let metadata_uuid = Uuid::from_u128(304);
let thread_id = ThreadId::from_string(&filename_uuid.to_string()).expect("valid thread id");
@@ -861,7 +828,7 @@ mod tests {
#[tokio::test]
async fn update_thread_metadata_rejects_multi_field_patch_without_partial_write() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(305);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
@@ -896,21 +863,12 @@ mod tests {
async fn update_thread_metadata_keeps_archived_thread_archived_in_sqlite() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(306);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-00-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
codex_rollout::state_db::reconcile_rollout(
Some(runtime.as_ref()),
archived_path.as_path(),
@@ -959,21 +917,12 @@ mod tests {
async fn update_thread_metadata_keeps_live_archived_thread_archived_in_sqlite() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(308);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-30-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
codex_rollout::state_db::reconcile_rollout(
Some(runtime.as_ref()),
archived_path.as_path(),