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

Make thread store process-scoped (#19474)

Browse Source 
- Build one app-server process ThreadStore from startup config and share
it with ThreadManager and CodexMessageProcessor.
- Remove per-thread/fork store reconstruction so effective thread config
cannot switch the persistence backend.
- Add params to ThreadStore create/resume for specifying thread
metadata, since otherwise the metadata from store creation would be used
(incorrectly).
This commit is contained in:
Tom
2026-04-30 21:24:59 -07:00
committed by GitHub
Unverified
parent f50c02d7bc
commit fe05acad23
55 changed files with 1076 additions and 514 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/thread-store/src/lib.rs
+2
View File
@@ -19,6 +19,7 @@ pub use in_memory::InMemoryThreadStoreCalls;
pub use live_thread::LiveThread;
pub use live_thread::LiveThreadInitGuard;
pub use local::LocalThreadStore;
pub use local::LocalThreadStoreConfig;
pub use remote::RemoteThreadStore;
pub use store::ThreadStore;
pub use types::AppendThreadItemsParams;
@@ -37,5 +38,6 @@ pub use types::StoredThreadHistory;
pub use types::ThreadEventPersistenceMode;
pub use types::ThreadMetadataPatch;
pub use types::ThreadPage;
pub use types::ThreadPersistenceMetadata;
pub use types::ThreadSortKey;
pub use types::UpdateThreadMetadataParams;
codex-rs/thread-store/src/local/archive_thread.rs
+4 -4
View File
@@ -48,7 +48,7 @@ pub(super) async fn archive_thread(
}
})?;
if let Some(ctx) = codex_rollout::state_db::get_state_db(&store.config).await {
if let Some(ctx) = store.state_db().await {
let _ = ctx
.mark_archived(thread_id, archived_path.as_path(), Utc::now())
.await;
@@ -130,7 +130,7 @@ mod tests {
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.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -144,10 +144,10 @@ mod tests {
Utc::now(),
SessionSource::Cli,
);
builder.model_provider = Some(config.model_provider_id.clone());
builder.model_provider = Some(config.default_model_provider_id.clone());
builder.cwd = home.path().to_path_buf();
builder.cli_version = Some("test_version".to_string());
let metadata = builder.build(config.model_provider_id.as_str());
let metadata = builder.build(config.default_model_provider_id.as_str());
runtime
.upsert_thread(&metadata)
.await
codex-rs/thread-store/src/local/create_thread.rs
+17 -1
View File
@@ -3,7 +3,9 @@ use crate::CreateThreadParams;
use crate::ThreadEventPersistenceMode;
use crate::ThreadStoreError;
use crate::ThreadStoreResult;
use codex_protocol::protocol::ThreadMemoryMode;
use codex_rollout::EventPersistenceMode;
use codex_rollout::RolloutConfig;
use codex_rollout::RolloutRecorder;
use codex_rollout::RolloutRecorderParams;
@@ -11,9 +13,23 @@ pub(super) async fn create_thread(
store: &LocalThreadStore,
params: CreateThreadParams,
) -> ThreadStoreResult<RolloutRecorder> {
let cwd = params
.metadata
.cwd
.clone()
.ok_or_else(|| ThreadStoreError::InvalidRequest {
message: "local thread store requires a cwd".to_string(),
})?;
let config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
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 recorder = RolloutRecorder::new(
&store.config,
&config,
RolloutRecorderParams::new(
params.thread_id,
params.forked_from_id,
codex-rs/thread-store/src/local/helpers.rs
+8
View File
@@ -13,6 +13,7 @@ use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::GitInfo;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::SessionSource;
use codex_rollout::ARCHIVED_SESSIONS_SUBDIR;
use codex_rollout::ThreadItem;
use codex_state::ThreadMetadata;
@@ -51,6 +52,13 @@ pub(super) fn scoped_rollout_path(
}
}
pub(super) fn rollout_path_is_archived(codex_home: &Path, path: &Path) -> bool {
path.starts_with(codex_home.join(ARCHIVED_SESSIONS_SUBDIR))
|| path
.components()
.any(|component| component.as_os_str() == OsStr::new(ARCHIVED_SESSIONS_SUBDIR))
}
pub(super) fn matching_rollout_file_name(
rollout_path: &Path,
thread_id: ThreadId,
codex-rs/thread-store/src/local/list_threads.rs
+18 -9
View File
@@ -39,8 +39,16 @@ pub(super) async fn list_threads(
SortDirection::Asc => codex_rollout::SortDirection::Asc,
SortDirection::Desc => codex_rollout::SortDirection::Desc,
};
let rollout_config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
cwd: store.config.codex_home.clone(),
model_provider_id: store.config.default_model_provider_id.clone(),
generate_memories: false,
};
let page = list_rollout_threads(
&store.config,
&rollout_config,
store.config.default_model_provider_id.as_str(),
&params,
cursor.as_ref(),
sort_key,
@@ -60,7 +68,7 @@ pub(super) async fn list_threads(
stored_thread_from_rollout_item(
item,
params.archived,
store.config.model_provider_id.as_str(),
store.config.default_model_provider_id.as_str(),
)
})
.collect::<Vec<_>>();
@@ -99,6 +107,7 @@ pub(super) async fn list_threads(
async fn list_rollout_threads(
config: &RolloutConfig,
default_model_provider_id: &str,
params: &ListThreadsParams,
cursor: Option<&codex_rollout::Cursor>,
sort_key: codex_rollout::ThreadSortKey,
@@ -114,7 +123,7 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
config.model_provider_id.as_str(),
default_model_provider_id,
params.search_term.as_deref(),
)
.await
@@ -128,7 +137,7 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
config.model_provider_id.as_str(),
default_model_provider_id,
params.search_term.as_deref(),
)
.await
@@ -142,7 +151,7 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
config.model_provider_id.as_str(),
default_model_provider_id,
params.search_term.as_deref(),
)
.await
@@ -156,7 +165,7 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
config.model_provider_id.as_str(),
default_model_provider_id,
params.search_term.as_deref(),
)
.await
@@ -230,7 +239,7 @@ mod tests {
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -245,10 +254,10 @@ mod tests {
created_at,
SessionSource::Cli,
);
builder.model_provider = Some(config.model_provider_id.clone());
builder.model_provider = Some(config.default_model_provider_id.clone());
builder.cwd = home.path().to_path_buf();
builder.cli_version = Some("test_version".to_string());
let mut metadata = builder.build(config.model_provider_id.as_str());
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
codex-rs/thread-store/src/local/live_writer.rs
+17 -1
View File
@@ -1,6 +1,8 @@
use std::path::PathBuf;
use codex_protocol::ThreadId;
use codex_protocol::protocol::ThreadMemoryMode;
use codex_rollout::RolloutConfig;
use codex_rollout::RolloutRecorder;
use codex_rollout::RolloutRecorderParams;
use codex_rollout::builder_from_items;
@@ -55,9 +57,23 @@ pub(super) async fn resume_thread(
let state_builder = history
.as_deref()
.and_then(|items| builder_from_items(items, rollout_path.as_path()));
let cwd = params
.metadata
.cwd
.clone()
.ok_or_else(|| ThreadStoreError::InvalidRequest {
message: "local thread store requires a cwd".to_string(),
})?;
let config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
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 recorder = RolloutRecorder::new(
&store.config,
&config,
RolloutRecorderParams::resume(
rollout_path,
create_thread::event_persistence_mode(params.event_persistence_mode),
codex-rs/thread-store/src/local/mod.rs
+194 -6
View File
@@ -12,7 +12,6 @@ mod test_support;
use async_trait::async_trait;
use codex_protocol::ThreadId;
use codex_rollout::RolloutConfig;
use codex_rollout::RolloutRecorder;
use codex_rollout::StateDbHandle;
use std::collections::HashMap;
@@ -41,11 +40,33 @@ use crate::UpdateThreadMetadataParams;
/// Local filesystem/SQLite-backed implementation of [`ThreadStore`].
#[derive(Clone)]
pub struct LocalThreadStore {
pub(super) config: RolloutConfig,
pub(super) config: LocalThreadStoreConfig,
live_recorders: Arc<Mutex<HashMap<ThreadId, RolloutRecorder>>>,
state_db: Arc<OnceCell<StateDbHandle>>,
}
/// Process-scoped configuration for local thread storage.
///
/// This describes where local storage lives. New-thread rollout metadata such
/// as cwd, provider, and memory mode is supplied when live persistence is opened.
#[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,
}
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(),
}
}
}
impl std::fmt::Debug for LocalThreadStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LocalThreadStore")
@@ -55,8 +76,8 @@ impl std::fmt::Debug for LocalThreadStore {
}
impl LocalThreadStore {
/// Create a local store from the rollout configuration used by existing local persistence.
pub fn new(config: RolloutConfig) -> Self {
/// Create a local store from process-scoped local storage configuration.
pub fn new(config: LocalThreadStoreConfig) -> Self {
Self {
config,
live_recorders: Arc::new(Mutex::new(HashMap::new())),
@@ -68,7 +89,13 @@ impl LocalThreadStore {
pub async fn state_db(&self) -> Option<StateDbHandle> {
self.state_db
.get_or_try_init(|| async {
codex_rollout::state_db::init(&self.config).await.ok_or(())
codex_rollout::state_db::init_with_roots(
self.config.codex_home.clone(),
self.config.sqlite_home.clone(),
self.config.default_model_provider_id.clone(),
)
.await
.ok_or(())
})
.await
.ok()
@@ -176,6 +203,16 @@ impl ThreadStore for LocalThreadStore {
params: LoadThreadHistoryParams,
) -> ThreadStoreResult<StoredThreadHistory> {
if let Ok(rollout_path) = live_writer::rollout_path(self, params.thread_id).await {
if !params.include_archived
&& helpers::rollout_path_is_archived(
self.config.codex_home.as_path(),
rollout_path.as_path(),
)
{
return Err(ThreadStoreError::InvalidRequest {
message: format!("thread {} is archived", params.thread_id),
});
}
return read_thread::read_thread_by_rollout_path(
self,
rollout_path,
@@ -251,11 +288,13 @@ mod tests {
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadMemoryMode;
use codex_protocol::protocol::UserMessageEvent;
use tempfile::TempDir;
use super::*;
use crate::ThreadEventPersistenceMode;
use crate::ThreadPersistenceMetadata;
use crate::local::test_support::test_config;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
@@ -309,6 +348,26 @@ 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()));
let thread_id = ThreadId::default();
let mut params = create_thread_params(thread_id);
params.metadata.cwd = None;
let err = store
.create_thread(params)
.await
.expect_err("local thread store should require cwd");
assert!(matches!(
err,
ThreadStoreError::InvalidRequest { message }
if message == "local thread store requires a cwd"
));
}
#[tokio::test]
async fn discard_thread_drops_unmaterialized_live_writer() {
let home = TempDir::new().expect("temp dir");
@@ -387,6 +446,7 @@ mod tests {
rollout_path: None,
history: None,
include_archived: true,
metadata: thread_metadata(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
@@ -427,6 +487,63 @@ mod tests {
assert!(err.to_string().contains("already has a live local writer"));
}
#[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()));
let thread_id = ThreadId::default();
store
.create_thread(create_thread_params(thread_id))
.await
.expect("create live thread");
let rollout_path = store
.live_rollout_path(thread_id)
.await
.expect("live rollout path");
let err = store
.resume_thread(ResumeThreadParams {
thread_id,
rollout_path: Some(rollout_path),
history: None,
include_archived: true,
metadata: thread_metadata(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
.expect_err("duplicate live resume should fail");
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
assert!(err.to_string().contains("already has a live local writer"));
}
#[tokio::test]
async fn resume_thread_rejects_missing_cwd() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()));
let uuid = uuid::Uuid::from_u128(407);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
write_session_file(home.path(), "2025-01-04T11-30-00", uuid).expect("session file");
let err = store
.resume_thread(ResumeThreadParams {
thread_id,
rollout_path: Some(rollout_path),
history: None,
include_archived: true,
metadata: ThreadPersistenceMetadata {
cwd: None,
model_provider: "test-provider".to_string(),
memory_mode: ThreadMemoryMode::Enabled,
},
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
.expect_err("missing cwd should fail");
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
assert!(err.to_string().contains("requires a cwd"));
}
#[tokio::test]
async fn load_history_uses_live_writer_rollout_path() {
let home = TempDir::new().expect("temp dir");
@@ -443,6 +560,7 @@ mod tests {
rollout_path: Some(rollout_path),
history: None,
include_archived: true,
metadata: thread_metadata(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
@@ -475,6 +593,46 @@ mod tests {
}));
}
#[tokio::test]
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()));
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)
.expect("external session file");
store
.resume_thread(ResumeThreadParams {
thread_id,
rollout_path: Some(rollout_path.clone()),
history: None,
include_archived: true,
metadata: thread_metadata(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
.expect("resume live thread");
let thread = store
.read_thread(ReadThreadParams {
thread_id,
include_archived: false,
include_history: true,
})
.await
.expect("read external live thread");
assert_eq!(thread.rollout_path, Some(rollout_path));
assert!(thread.history.expect("history").items.iter().any(|item| {
matches!(
item,
RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == "Hello from user"
)
}));
}
#[tokio::test]
async fn load_history_uses_live_writer_rollout_path_for_archived_source() {
let home = TempDir::new().expect("temp dir");
@@ -490,6 +648,7 @@ mod tests {
rollout_path: Some(rollout_path),
history: None,
include_archived: true,
metadata: thread_metadata(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
@@ -506,12 +665,32 @@ mod tests {
.await
.expect("flush live thread");
let history = store
let err = store
.read_thread(ReadThreadParams {
thread_id,
include_archived: false,
include_history: false,
})
.await
.expect_err("active-only read should reject archived live thread");
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
let err = store
.load_history(LoadThreadHistoryParams {
thread_id,
include_archived: false,
})
.await
.expect_err("active-only history should reject archived live thread");
assert!(matches!(err, ThreadStoreError::InvalidRequest { .. }));
assert!(err.to_string().contains("archived"));
let history = store
.load_history(LoadThreadHistoryParams {
thread_id,
include_archived: true,
})
.await
.expect("load archived live history");
assert!(history.items.iter().any(|item| {
@@ -574,10 +753,19 @@ mod tests {
source: SessionSource::Exec,
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
metadata: thread_metadata(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
}
}
fn thread_metadata() -> ThreadPersistenceMetadata {
ThreadPersistenceMetadata {
cwd: Some(std::env::current_dir().expect("cwd")),
model_provider: "test-provider".to_string(),
memory_mode: ThreadMemoryMode::Enabled,
}
}
fn user_message_item(message: &str) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
message: message.to_string(),
codex-rs/thread-store/src/local/read_thread.rs
+67 -48
View File
@@ -16,8 +16,10 @@ use codex_state::ThreadMetadata;
use super::LocalThreadStore;
use super::helpers::distinct_thread_metadata_title;
use super::helpers::git_info_from_parts;
use super::helpers::rollout_path_is_archived;
use super::helpers::set_thread_name_from_title;
use super::helpers::stored_thread_from_rollout_item;
use super::live_writer;
use crate::ReadThreadParams;
use crate::StoredThread;
use crate::StoredThreadHistory;
@@ -30,7 +32,12 @@ pub(super) async fn read_thread(
) -> ThreadStoreResult<StoredThread> {
let thread_id = params.thread_id;
if let Some(metadata) = read_sqlite_metadata(store, thread_id).await
&& (params.include_archived || metadata.archived_at.is_none())
&& (params.include_archived
|| (metadata.archived_at.is_none()
&& !rollout_path_is_archived(
store.config.codex_home.as_path(),
metadata.rollout_path.as_path(),
)))
&& (!params.include_history
|| sqlite_rollout_path_can_load_history_for_thread(
store,
@@ -44,6 +51,7 @@ pub(super) async fn read_thread(
&& let Some(rollout_path) = thread.rollout_path.clone()
&& let Ok(mut rollout_thread) = read_thread_from_rollout_path(store, rollout_path).await
&& rollout_thread.thread_id == thread_id
&& (params.include_archived || rollout_thread.archived_at.is_none())
&& !rollout_thread.preview.is_empty()
{
if thread.name.is_some() {
@@ -153,6 +161,17 @@ async fn resolve_rollout_path(
thread_id: codex_protocol::ThreadId,
include_archived: bool,
) -> ThreadStoreResult<Option<std::path::PathBuf>> {
if let Ok(path) = live_writer::rollout_path(store, thread_id).await
&& tokio::fs::try_exists(path.as_path()).await.map_err(|err| {
ThreadStoreError::InvalidRequest {
message: format!("failed to check rollout path for thread id {thread_id}: {err}"),
}
})?
&& (include_archived || !rollout_path_is_archived(store.config.codex_home.as_path(), &path))
{
return Ok(Some(path));
}
if include_archived {
match find_thread_path_by_id_str(store.config.codex_home.as_path(), &thread_id.to_string())
.await
@@ -185,21 +204,25 @@ async fn read_thread_from_rollout_path(
let Some(item) = read_thread_item_from_rollout(path.clone()).await else {
return stored_thread_from_session_meta(store, path).await;
};
let archived = path.starts_with(
store
.config
.codex_home
.join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR),
);
let mut thread =
stored_thread_from_rollout_item(item, archived, store.config.model_provider_id.as_str())
.ok_or_else(|| ThreadStoreError::Internal {
message: format!("failed to read thread id from {}", path.display()),
})?;
thread.forked_from_id = read_session_meta_line(path.as_path())
.await
.ok()
.and_then(|meta_line| meta_line.meta.forked_from_id);
let archived = rollout_path_is_archived(store.config.codex_home.as_path(), path.as_path());
let mut thread = stored_thread_from_rollout_item(
item,
archived,
store.config.default_model_provider_id.as_str(),
)
.ok_or_else(|| ThreadStoreError::Internal {
message: format!("failed to read thread id from {}", path.display()),
})?;
if let Ok(meta_line) = read_session_meta_line(path.as_path()).await {
thread.forked_from_id = meta_line.meta.forked_from_id;
if let Some(model_provider) = meta_line
.meta
.model_provider
.filter(|provider| !provider.is_empty())
{
thread.model_provider = model_provider;
}
}
if let Ok(Some(title)) =
find_thread_name_by_id(store.config.codex_home.as_path(), &thread.thread_id).await
{
@@ -225,7 +248,7 @@ async fn read_sqlite_metadata(
) -> Option<ThreadMetadata> {
let runtime = StateRuntime::init(
store.config.sqlite_home.clone(),
store.config.model_provider_id.clone(),
store.config.default_model_provider_id.clone(),
)
.await
.ok()?;
@@ -254,7 +277,7 @@ async fn stored_thread_from_sqlite_metadata(
preview: metadata.first_user_message.clone().unwrap_or_default(),
name,
model_provider: if metadata.model_provider.is_empty() {
store.config.model_provider_id.clone()
store.config.default_model_provider_id.clone()
} else {
metadata.model_provider
},
@@ -294,12 +317,7 @@ async fn stored_thread_from_session_meta(
.map_err(|err| ThreadStoreError::Internal {
message: format!("failed to read thread {}: {err}", path.display()),
})?;
let archived = path.starts_with(
store
.config
.codex_home
.join(codex_rollout::ARCHIVED_SESSIONS_SUBDIR),
);
let archived = rollout_path_is_archived(store.config.codex_home.as_path(), path.as_path());
Ok(stored_thread_from_meta_line(
store, meta_line, path, archived,
))
@@ -327,7 +345,7 @@ fn stored_thread_from_meta_line(
.meta
.model_provider
.filter(|provider| !provider.is_empty())
.unwrap_or_else(|| store.config.model_provider_id.clone()),
.unwrap_or_else(|| store.config.default_model_provider_id.clone()),
model: None,
reasoning_effort: None,
created_at,
@@ -459,7 +477,7 @@ mod tests {
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.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -469,10 +487,10 @@ mod tests {
Utc::now(),
SessionSource::Cli,
);
builder.model_provider = Some(config.model_provider_id.clone());
builder.model_provider = Some(config.default_model_provider_id.clone());
builder.git_branch = Some("sqlite-branch".to_string());
runtime
.upsert_thread(&builder.build(config.model_provider_id.as_str()))
.upsert_thread(&builder.build(config.default_model_provider_id.as_str()))
.await
.expect("state db upsert should succeed");
@@ -606,16 +624,16 @@ mod tests {
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.model_provider_id.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);
builder.model_provider = Some(config.model_provider_id.clone());
builder.model_provider = Some(config.default_model_provider_id.clone());
builder.cwd = home.path().to_path_buf();
builder.cli_version = Some("test_version".to_string());
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.title = "Saved title".to_string();
metadata.first_user_message = Some("Hello from user".to_string());
runtime
@@ -674,7 +692,7 @@ mod tests {
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -684,9 +702,9 @@ mod tests {
Utc::now(),
SessionSource::Cli,
);
builder.model_provider = Some(config.model_provider_id.clone());
builder.model_provider = Some(config.default_model_provider_id.clone());
builder.cwd = home.path().join("sqlite-workspace");
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.title = "Saved title".to_string();
metadata.first_user_message = Some("Hello from sqlite".to_string());
runtime
@@ -707,6 +725,7 @@ mod tests {
assert_eq!(thread.rollout_path, Some(rollout_path));
assert_eq!(thread.preview, "Hello from rollout");
assert_eq!(thread.name, Some("Saved title".to_string()));
assert_eq!(thread.model_provider, "rollout-provider");
assert_eq!(thread.cwd, rollout_cwd);
}
@@ -761,7 +780,7 @@ mod tests {
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -774,7 +793,7 @@ mod tests {
builder.model_provider = Some("sqlite-provider".to_string());
builder.cwd = home.path().join("workspace");
builder.cli_version = Some("sqlite-cli".to_string());
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.title = "Command-only thread".to_string();
runtime
.upsert_thread(&metadata)
@@ -815,7 +834,7 @@ mod tests {
let stale_path = external.path().join("missing-rollout.jsonl");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -826,7 +845,7 @@ mod tests {
SessionSource::Cli,
);
builder.model_provider = Some("stale-sqlite-provider".to_string());
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.first_user_message = Some("stale sqlite preview".to_string());
runtime
.upsert_thread(&metadata)
@@ -845,7 +864,7 @@ mod tests {
assert_eq!(thread.thread_id, thread_id);
assert_eq!(thread.rollout_path, Some(rollout_path));
assert_eq!(thread.preview, "Hello from user");
assert_eq!(thread.model_provider, config.model_provider_id);
assert_eq!(thread.model_provider, config.default_model_provider_id);
let history = thread.history.expect("history should load");
assert_eq!(history.thread_id, thread_id);
assert_eq!(history.items.len(), 2);
@@ -866,14 +885,14 @@ mod tests {
.expect("other session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let mut builder =
ThreadMetadataBuilder::new(thread_id, stale_path, Utc::now(), SessionSource::Cli);
builder.model_provider = Some("wrong-sqlite-provider".to_string());
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.first_user_message = Some("wrong sqlite preview".to_string());
runtime
.upsert_thread(&metadata)
@@ -892,7 +911,7 @@ mod tests {
assert_eq!(thread.thread_id, thread_id);
assert_eq!(thread.rollout_path, Some(rollout_path));
assert_eq!(thread.preview, "Hello from user");
assert_eq!(thread.model_provider, config.model_provider_id);
assert_eq!(thread.model_provider, config.default_model_provider_id);
let history = thread.history.expect("history should load");
assert_eq!(history.thread_id, thread_id);
assert_eq!(history.items.len(), 2);
@@ -964,7 +983,7 @@ mod tests {
.join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl"));
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -977,7 +996,7 @@ mod tests {
builder.model_provider = Some("sqlite-provider".to_string());
builder.cwd = external.path().join("workspace");
builder.cli_version = Some("sqlite-cli".to_string());
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.title = "SQLite title".to_string();
metadata.first_user_message = Some("SQLite preview".to_string());
metadata.model = Some("sqlite-model".to_string());
@@ -1022,14 +1041,14 @@ mod tests {
.join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl"));
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.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);
builder.archived_at = Some(Utc::now());
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.first_user_message = Some("Archived SQLite preview".to_string());
runtime
.upsert_thread(&metadata)
@@ -1077,7 +1096,7 @@ mod tests {
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -1088,7 +1107,7 @@ mod tests {
SessionSource::Cli,
);
builder.archived_at = Some(Utc::now());
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.first_user_message = Some("Archived SQLite preview".to_string());
runtime
.upsert_thread(&metadata)
codex-rs/thread-store/src/local/test_support.rs
+5 -6
View File
@@ -4,16 +4,15 @@ use std::path::Path;
use std::path::PathBuf;
use codex_rollout::ARCHIVED_SESSIONS_SUBDIR;
use codex_rollout::RolloutConfig;
use uuid::Uuid;
pub(super) fn test_config(codex_home: &Path) -> RolloutConfig {
RolloutConfig {
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(),
cwd: codex_home.to_path_buf(),
model_provider_id: "test-provider".to_string(),
generate_memories: true,
default_model_provider_id: "test-provider".to_string(),
}
}
codex-rs/thread-store/src/local/unarchive_thread.rs
+5 -5
View File
@@ -71,7 +71,7 @@ pub(super) async fn unarchive_thread(
message: format!("failed to update unarchived thread timestamp: {err}"),
})?;
if let Some(ctx) = codex_rollout::state_db::get_state_db(&store.config).await {
if let Some(ctx) = store.state_db().await {
let _ = ctx
.mark_unarchived(thread_id, restored_path.as_path())
.await;
@@ -88,7 +88,7 @@ pub(super) async fn unarchive_thread(
stored_thread_from_rollout_item(
item,
/*archived*/ false,
store.config.model_provider_id.as_str(),
store.config.default_model_provider_id.as_str(),
)
.ok_or_else(|| ThreadStoreError::Internal {
message: format!(
@@ -154,7 +154,7 @@ mod tests {
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -168,10 +168,10 @@ mod tests {
Utc::now(),
SessionSource::Cli,
);
builder.model_provider = Some(config.model_provider_id.clone());
builder.model_provider = Some(config.default_model_provider_id.clone());
builder.cwd = home.path().to_path_buf();
builder.cli_version = Some("test_version".to_string());
let mut metadata = builder.build(config.model_provider_id.as_str());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.archived_at = Some(metadata.updated_at);
runtime
.upsert_thread(&metadata)
codex-rs/thread-store/src/local/update_thread_metadata.rs
+17 -6
View File
@@ -58,7 +58,7 @@ pub(super) async fn update_thread_metadata(
codex_rollout::state_db::reconcile_rollout(
state_db_ctx.as_deref(),
resolved_rollout_path.path.as_path(),
store.config.model_provider_id.as_str(),
store.config.default_model_provider_id.as_str(),
/*builder*/ None,
&[],
/*archived_only*/ resolved_rollout_path.archived.then_some(true),
@@ -203,6 +203,7 @@ mod tests {
use crate::ResumeThreadParams;
use crate::ThreadEventPersistenceMode;
use crate::ThreadMetadataPatch;
use crate::ThreadPersistenceMetadata;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::test_config;
@@ -254,7 +255,7 @@ mod tests {
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.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -299,6 +300,7 @@ mod tests {
rollout_path: Some(path.clone()),
history: None,
include_archived: true,
metadata: test_thread_metadata(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
@@ -400,7 +402,7 @@ mod tests {
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -411,7 +413,7 @@ mod tests {
codex_rollout::state_db::reconcile_rollout(
Some(runtime.as_ref()),
archived_path.as_path(),
config.model_provider_id.as_str(),
config.default_model_provider_id.as_str(),
/*builder*/ None,
&[],
/*archived_only*/ Some(true),
@@ -463,7 +465,7 @@ mod tests {
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.model_provider_id.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
@@ -474,7 +476,7 @@ mod tests {
codex_rollout::state_db::reconcile_rollout(
Some(runtime.as_ref()),
archived_path.as_path(),
config.model_provider_id.as_str(),
config.default_model_provider_id.as_str(),
/*builder*/ None,
&[],
/*archived_only*/ Some(true),
@@ -487,6 +489,7 @@ mod tests {
rollout_path: Some(archived_path.clone()),
history: None,
include_archived: true,
metadata: test_thread_metadata(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
@@ -516,6 +519,14 @@ mod tests {
);
}
fn test_thread_metadata() -> ThreadPersistenceMetadata {
ThreadPersistenceMetadata {
cwd: Some(std::env::current_dir().expect("cwd")),
model_provider: "test-provider".to_string(),
memory_mode: ThreadMemoryMode::Enabled,
}
}
fn last_rollout_item(path: &std::path::Path) -> Value {
let last_line = std::fs::read_to_string(path)
.expect("read rollout")
codex-rs/thread-store/src/remote/helpers.rs
+7
View File
@@ -25,6 +25,7 @@ use crate::StoredThread;
use crate::StoredThreadHistory;
use crate::ThreadEventPersistenceMode;
use crate::ThreadMetadataPatch;
use crate::ThreadPersistenceMetadata;
use crate::ThreadSortKey;
use crate::ThreadStoreError;
use crate::ThreadStoreResult;
@@ -186,6 +187,12 @@ pub(super) fn dynamic_tools_json(
serialize_json_vec(dynamic_tools, "dynamic_tool")
}
pub(super) fn thread_persistence_metadata_json(
metadata: &ThreadPersistenceMetadata,
) -> ThreadStoreResult<String> {
serialize_json(metadata, "thread_persistence_metadata")
}
pub(super) fn rollout_items_json(items: &[RolloutItem]) -> ThreadStoreResult<Vec<String>> {
serialize_json_vec(items, "rollout_item")
}
codex-rs/thread-store/src/remote/mod.rs
+147
View File
@@ -69,6 +69,7 @@ impl ThreadStore for RemoteThreadStore {
params.event_persistence_mode,
)
.into(),
metadata_json: helpers::thread_persistence_metadata_json(&params.metadata)?,
};
self.client()
.await?
@@ -96,6 +97,7 @@ impl ThreadStore for RemoteThreadStore {
params.event_persistence_mode,
)
.into(),
metadata_json: helpers::thread_persistence_metadata_json(&params.metadata)?,
};
self.client()
.await?
@@ -260,3 +262,148 @@ impl ThreadStore for RemoteThreadStore {
helpers::stored_thread_from_proto(thread)
}
}
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use codex_protocol::ThreadId;
use codex_protocol::models::BaseInstructions;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadMemoryMode;
use pretty_assertions::assert_eq;
use tokio::sync::mpsc;
use tonic::Request;
use tonic::Response;
use tonic::Status;
use tonic::transport::Server;
use super::*;
use crate::ThreadEventPersistenceMode;
use crate::ThreadPersistenceMetadata;
use proto::thread_store_server;
use proto::thread_store_server::ThreadStoreServer;
enum RecordedRequest {
Create(proto::CreateThreadRequest),
Resume(proto::ResumeThreadRequest),
}
struct TestServer {
requests_tx: mpsc::UnboundedSender<RecordedRequest>,
}
#[tonic::async_trait]
impl thread_store_server::ThreadStore for TestServer {
async fn create_thread(
&self,
request: Request<proto::CreateThreadRequest>,
) -> Result<Response<proto::Empty>, Status> {
self.requests_tx
.send(RecordedRequest::Create(request.into_inner()))
.expect("record create request");
Ok(Response::new(proto::Empty {}))
}
async fn resume_thread(
&self,
request: Request<proto::ResumeThreadRequest>,
) -> Result<Response<proto::Empty>, Status> {
self.requests_tx
.send(RecordedRequest::Resume(request.into_inner()))
.expect("record resume request");
Ok(Response::new(proto::Empty {}))
}
async fn list_threads(
&self,
_request: Request<proto::ListThreadsRequest>,
) -> Result<Response<proto::ListThreadsResponse>, Status> {
Err(Status::unimplemented("not implemented"))
}
}
async fn test_store() -> (RemoteThreadStore, mpsc::UnboundedReceiver<RecordedRequest>) {
let (requests_tx, requests_rx) = mpsc::unbounded_channel();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind test server");
let addr = listener.local_addr().expect("test server addr");
tokio::spawn(async move {
Server::builder()
.add_service(ThreadStoreServer::new(TestServer { requests_tx }))
.serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
.await
.expect("test server");
});
(
RemoteThreadStore::new(format!("http://{addr}")),
requests_rx,
)
}
#[tokio::test]
async fn create_thread_forwards_metadata() {
let (store, mut requests_rx) = test_store().await;
let metadata = ThreadPersistenceMetadata {
cwd: Some(PathBuf::from("/workspace")),
model_provider: "test-provider".to_string(),
memory_mode: ThreadMemoryMode::Enabled,
};
store
.create_thread(CreateThreadParams {
thread_id: ThreadId::new(),
forked_from_id: None,
source: SessionSource::Exec,
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
metadata: metadata.clone(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
.expect("create thread");
let Some(RecordedRequest::Create(request)) = requests_rx.recv().await else {
panic!("expected create request");
};
assert_eq!(
serde_json::from_str::<ThreadPersistenceMetadata>(&request.metadata_json)
.expect("metadata json"),
metadata
);
}
#[tokio::test]
async fn resume_thread_forwards_metadata() {
let (store, mut requests_rx) = test_store().await;
let metadata = ThreadPersistenceMetadata {
cwd: Some(PathBuf::from("/workspace")),
model_provider: "test-provider".to_string(),
memory_mode: ThreadMemoryMode::Disabled,
};
store
.resume_thread(ResumeThreadParams {
thread_id: ThreadId::new(),
rollout_path: None,
history: None,
include_archived: false,
metadata: metadata.clone(),
event_persistence_mode: ThreadEventPersistenceMode::Limited,
})
.await
.expect("resume thread");
let Some(RecordedRequest::Resume(request)) = requests_rx.recv().await else {
panic!("expected resume request");
};
assert_eq!(
serde_json::from_str::<ThreadPersistenceMetadata>(&request.metadata_json)
.expect("metadata json"),
metadata
);
}
}
codex-rs/thread-store/src/remote/proto/codex.thread_store.v1.proto
+2
View File
@@ -31,6 +31,7 @@ message CreateThreadRequest {
string base_instructions_json = 4;
repeated string dynamic_tools_json = 5;
ThreadEventPersistenceMode event_persistence_mode = 6;
string metadata_json = 7;
}
message ResumeThreadRequest {
@@ -40,6 +41,7 @@ message ResumeThreadRequest {
bool has_history = 4;
bool include_archived = 5;
ThreadEventPersistenceMode event_persistence_mode = 6;
string metadata_json = 7;
}
message AppendThreadItemsRequest {
codex-rs/thread-store/src/remote/proto/codex.thread_store.v1.rs
+4
View File
@@ -22,6 +22,8 @@ pub struct CreateThreadRequest {
pub dynamic_tools_json: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
#[prost(enumeration = "ThreadEventPersistenceMode", tag = "6")]
pub event_persistence_mode: i32,
#[prost(string, tag = "7")]
pub metadata_json: ::prost::alloc::string::String,
}
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
pub struct ResumeThreadRequest {
@@ -37,6 +39,8 @@ pub struct ResumeThreadRequest {
pub include_archived: bool,
#[prost(enumeration = "ThreadEventPersistenceMode", tag = "6")]
pub event_persistence_mode: i32,
#[prost(string, tag = "7")]
pub metadata_json: ::prost::alloc::string::String,
}
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
pub struct AppendThreadItemsRequest {
codex-rs/thread-store/src/types.rs
+17
View File
@@ -26,6 +26,19 @@ pub enum ThreadEventPersistenceMode {
Extended,
}
/// Thread-scoped metadata used when opening live persistence.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ThreadPersistenceMetadata {
/// Effective working directory for environment-backed threads.
///
/// `None` means the thread has no filesystem/environment context.
pub cwd: Option<PathBuf>,
/// Model provider associated with the thread.
pub model_provider: String,
/// Memory mode associated with the live thread.
pub memory_mode: MemoryMode,
}
/// Parameters required to create a persisted thread.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CreateThreadParams {
@@ -39,6 +52,8 @@ pub struct CreateThreadParams {
pub base_instructions: BaseInstructions,
/// Dynamic tools available to the thread at startup.
pub dynamic_tools: Vec<DynamicToolSpec>,
/// Metadata captured for the newly created thread.
pub metadata: ThreadPersistenceMetadata,
/// Whether persistence should include the extended event surface.
pub event_persistence_mode: ThreadEventPersistenceMode,
}
@@ -54,6 +69,8 @@ pub struct ResumeThreadParams {
pub history: Option<Vec<RolloutItem>>,
/// Whether archived threads may be reopened.
pub include_archived: bool,
/// Metadata for future writes appended to the resumed live thread.
pub metadata: ThreadPersistenceMetadata,
/// Whether persistence should include the extended event surface.
pub event_persistence_mode: ThreadEventPersistenceMode,
}