mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
[codex] Add local thread store listing (#17824)
Builds on top of #17659 Move the filesystem + sqlite thread listing-related operations inside of a local ThreadStore implementation and call ThreadStore from the places that used to perform these filesystem/sqlite operations. This is the first of a series of PRs that will implement the rest of the local ThreadStore. Testing: - added unit tests for the thread store implementation - adjusted some unit tests in the realtime + personality packages whose callsites changed. Specifically I'm trying to hide ThreadMetadata inside of the local implementation and make ThreadMetadata a sqlite implementation detail concern rather than a public interface, preferring the more generate StoredThread interface instead - added a corner case test for the personality migration package that wasn't covered by the existing test suite - adjust the behavior of searched thread listing to run the existing local rollout repair/backfill pass _before_ querying SQLite results, so callers using ThreadStore::list_threads do not miss matches after a partial metadata warm-up
This commit is contained in:
@@ -136,9 +136,11 @@ use codex_protocol::request_permissions::RequestPermissionsResponse;
|
||||
use codex_protocol::request_user_input::RequestUserInputArgs;
|
||||
use codex_protocol::request_user_input::RequestUserInputResponse;
|
||||
use codex_rmcp_client::ElicitationResponse;
|
||||
use codex_rollout::RolloutConfig;
|
||||
use codex_rollout::state_db;
|
||||
use codex_shell_command::parse_command::parse_command;
|
||||
use codex_terminal_detection::user_agent;
|
||||
use codex_thread_store::LocalThreadStore;
|
||||
use codex_tools::filter_tool_suggest_discoverable_tools_for_client;
|
||||
use codex_utils_output_truncation::TruncationPolicy;
|
||||
use codex_utils_stream_parser::AssistantTextChunk;
|
||||
@@ -2127,6 +2129,7 @@ impl Session {
|
||||
network_proxy,
|
||||
network_approval: Arc::clone(&network_approval),
|
||||
state_db: state_db_ctx.clone(),
|
||||
thread_store: LocalThreadStore::new(RolloutConfig::from_view(config.as_ref())),
|
||||
model_client: ModelClient::new(
|
||||
Some(Arc::clone(&auth_manager)),
|
||||
conversation_id,
|
||||
|
||||
@@ -2836,6 +2836,9 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
|
||||
network_proxy: None,
|
||||
network_approval: Arc::clone(&network_approval),
|
||||
state_db: None,
|
||||
thread_store: codex_thread_store::LocalThreadStore::new(
|
||||
codex_rollout::RolloutConfig::from_view(config.as_ref()),
|
||||
),
|
||||
model_client: ModelClient::new(
|
||||
Some(auth_manager.clone()),
|
||||
conversation_id,
|
||||
@@ -3693,6 +3696,9 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
|
||||
network_proxy: None,
|
||||
network_approval: Arc::clone(&network_approval),
|
||||
state_db: None,
|
||||
thread_store: codex_thread_store::LocalThreadStore::new(
|
||||
codex_rollout::RolloutConfig::from_view(config.as_ref()),
|
||||
),
|
||||
model_client: ModelClient::new(
|
||||
Some(Arc::clone(&auth_manager)),
|
||||
conversation_id,
|
||||
|
||||
@@ -1,14 +1,10 @@
|
||||
use crate::config::edit::ConfigEditsBuilder;
|
||||
use crate::rollout::ARCHIVED_SESSIONS_SUBDIR;
|
||||
use crate::rollout::SESSIONS_SUBDIR;
|
||||
use crate::rollout::list::ThreadListConfig;
|
||||
use crate::rollout::list::ThreadListLayout;
|
||||
use crate::rollout::list::ThreadSortKey;
|
||||
use crate::rollout::list::get_threads_in_root;
|
||||
use codex_config::config_toml::ConfigToml;
|
||||
use codex_protocol::config_types::Personality;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_rollout::state_db;
|
||||
use codex_thread_store::ListThreadsParams;
|
||||
use codex_thread_store::LocalThreadStore;
|
||||
use codex_thread_store::ThreadSortKey;
|
||||
use codex_thread_store::ThreadStore;
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
use tokio::fs::OpenOptions;
|
||||
@@ -64,57 +60,33 @@ pub async fn maybe_migrate_personality(
|
||||
}
|
||||
|
||||
async fn has_recorded_sessions(codex_home: &Path, default_provider: &str) -> io::Result<bool> {
|
||||
let allowed_sources: &[SessionSource] = &[];
|
||||
let store = LocalThreadStore::new(codex_rollout::RolloutConfig {
|
||||
codex_home: codex_home.to_path_buf(),
|
||||
sqlite_home: codex_home.to_path_buf(),
|
||||
cwd: codex_home.to_path_buf(),
|
||||
model_provider_id: default_provider.to_string(),
|
||||
generate_memories: false,
|
||||
});
|
||||
if has_threads(&store, /*archived*/ false).await? {
|
||||
return Ok(true);
|
||||
}
|
||||
has_threads(&store, /*archived*/ true).await
|
||||
}
|
||||
|
||||
if let Some(state_db_ctx) = state_db::open_if_present(codex_home, default_provider).await
|
||||
&& let Some(ids) = state_db::list_thread_ids_db(
|
||||
Some(state_db_ctx.as_ref()),
|
||||
codex_home,
|
||||
/*page_size*/ 1,
|
||||
/*cursor*/ None,
|
||||
ThreadSortKey::CreatedAt,
|
||||
allowed_sources,
|
||||
/*model_providers*/ None,
|
||||
/*archived_only*/ false,
|
||||
"personality_migration",
|
||||
)
|
||||
async fn has_threads(store: &LocalThreadStore, archived: bool) -> io::Result<bool> {
|
||||
store
|
||||
.list_threads(ListThreadsParams {
|
||||
page_size: 1,
|
||||
cursor: None,
|
||||
sort_key: ThreadSortKey::CreatedAt,
|
||||
allowed_sources: Vec::new(),
|
||||
model_providers: None,
|
||||
archived,
|
||||
search_term: None,
|
||||
})
|
||||
.await
|
||||
&& !ids.is_empty()
|
||||
{
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let sessions = get_threads_in_root(
|
||||
codex_home.join(SESSIONS_SUBDIR),
|
||||
/*page_size*/ 1,
|
||||
/*cursor*/ None,
|
||||
ThreadSortKey::CreatedAt,
|
||||
ThreadListConfig {
|
||||
allowed_sources,
|
||||
model_providers: None,
|
||||
default_provider,
|
||||
layout: ThreadListLayout::NestedByDate,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
if !sessions.items.is_empty() {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let archived_sessions = get_threads_in_root(
|
||||
codex_home.join(ARCHIVED_SESSIONS_SUBDIR),
|
||||
/*page_size*/ 1,
|
||||
/*cursor*/ None,
|
||||
ThreadSortKey::CreatedAt,
|
||||
ThreadListConfig {
|
||||
allowed_sources,
|
||||
model_providers: None,
|
||||
default_provider,
|
||||
layout: ThreadListLayout::Flat,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
Ok(!archived_sessions.items.is_empty())
|
||||
.map(|page| !page.items.is_empty())
|
||||
.map_err(io::Error::other)
|
||||
}
|
||||
|
||||
async fn create_marker(marker_path: &Path) -> io::Result<()> {
|
||||
|
||||
@@ -7,6 +7,8 @@ use codex_protocol::protocol::SessionMeta;
|
||||
use codex_protocol::protocol::SessionMetaLine;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::UserMessageEvent;
|
||||
use codex_rollout::ARCHIVED_SESSIONS_SUBDIR;
|
||||
use codex_rollout::SESSIONS_SUBDIR;
|
||||
use pretty_assertions::assert_eq;
|
||||
use tempfile::TempDir;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
@@ -25,6 +27,16 @@ async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> {
|
||||
.join("2025")
|
||||
.join("01")
|
||||
.join("01");
|
||||
write_rollout_with_user_event(&dir, thread_id).await
|
||||
}
|
||||
|
||||
async fn write_archived_session_with_user_event(codex_home: &Path) -> io::Result<()> {
|
||||
let thread_id = ThreadId::new();
|
||||
let dir = codex_home.join(ARCHIVED_SESSIONS_SUBDIR);
|
||||
write_rollout_with_user_event(&dir, thread_id).await
|
||||
}
|
||||
|
||||
async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::Result<()> {
|
||||
tokio::fs::create_dir_all(&dir).await?;
|
||||
let file_path = dir.join(format!("rollout-{TEST_TIMESTAMP}-{thread_id}.jsonl"));
|
||||
let mut file = tokio::fs::File::create(&file_path).await?;
|
||||
@@ -85,6 +97,22 @@ async fn applies_when_sessions_exist_and_no_personality() -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn applies_when_only_archived_sessions_exist_and_no_personality() -> io::Result<()> {
|
||||
let temp = TempDir::new()?;
|
||||
write_archived_session_with_user_event(temp.path()).await?;
|
||||
|
||||
let config_toml = ConfigToml::default();
|
||||
let status = maybe_migrate_personality(temp.path(), &config_toml).await?;
|
||||
|
||||
assert_eq!(status, PersonalityMigrationStatus::Applied);
|
||||
assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists());
|
||||
|
||||
let persisted = read_config_toml(temp.path()).await?;
|
||||
assert_eq!(persisted.personality, Some(Personality::Pragmatic));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_when_marker_exists() -> io::Result<()> {
|
||||
let temp = TempDir::new()?;
|
||||
|
||||
@@ -4,8 +4,10 @@ use crate::event_mapping::is_contextual_user_message_content;
|
||||
use chrono::Utc;
|
||||
use codex_git_utils::resolve_root_git_project_for_trust;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_state::SortKey;
|
||||
use codex_state::ThreadMetadata;
|
||||
use codex_thread_store::ListThreadsParams;
|
||||
use codex_thread_store::StoredThread;
|
||||
use codex_thread_store::ThreadSortKey;
|
||||
use codex_thread_store::ThreadStore;
|
||||
use codex_utils_output_truncation::TruncationPolicy;
|
||||
use codex_utils_output_truncation::truncate_text;
|
||||
use dirs::home_dir;
|
||||
@@ -98,7 +100,7 @@ pub(crate) async fn build_realtime_startup_context(
|
||||
}
|
||||
if let Some(section) = format_section(
|
||||
"Notes",
|
||||
Some("Built at realtime startup from the current thread history, persisted thread metadata in the state DB, and a bounded local workspace scan. This excludes repo memory instructions, AGENTS files, project-doc prompt blends, and memory summaries.".to_string()),
|
||||
Some("Built at realtime startup from the current thread history, local thread metadata, and a bounded local workspace scan. This excludes repo memory instructions, AGENTS files, project-doc prompt blends, and memory summaries.".to_string()),
|
||||
NOTES_SECTION_TOKEN_BUDGET,
|
||||
) {
|
||||
parts.push(section);
|
||||
@@ -117,33 +119,31 @@ pub(crate) async fn build_realtime_startup_context(
|
||||
Some(context)
|
||||
}
|
||||
|
||||
async fn load_recent_threads(sess: &Session) -> Vec<ThreadMetadata> {
|
||||
let Some(state_db) = sess.services.state_db.as_ref() else {
|
||||
return Vec::new();
|
||||
};
|
||||
|
||||
match state_db
|
||||
.list_threads(
|
||||
MAX_RECENT_THREADS,
|
||||
/*anchor*/ None,
|
||||
SortKey::UpdatedAt,
|
||||
&[],
|
||||
/*model_providers*/ None,
|
||||
/*archived_only*/ false,
|
||||
/*search_term*/ None,
|
||||
)
|
||||
async fn load_recent_threads(sess: &Session) -> Vec<StoredThread> {
|
||||
match sess
|
||||
.services
|
||||
.thread_store
|
||||
.list_threads(ListThreadsParams {
|
||||
page_size: MAX_RECENT_THREADS,
|
||||
cursor: None,
|
||||
sort_key: ThreadSortKey::UpdatedAt,
|
||||
allowed_sources: Vec::new(),
|
||||
model_providers: None,
|
||||
archived: false,
|
||||
search_term: None,
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(page) => page.items,
|
||||
Err(err) => {
|
||||
warn!("failed to load realtime startup threads from state db: {err}");
|
||||
warn!("failed to load realtime startup threads from thread store: {err}");
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_recent_work_section(cwd: &Path, recent_threads: &[ThreadMetadata]) -> Option<String> {
|
||||
let mut groups: HashMap<PathBuf, Vec<&ThreadMetadata>> = HashMap::new();
|
||||
fn build_recent_work_section(cwd: &Path, recent_threads: &[StoredThread]) -> Option<String> {
|
||||
let mut groups: HashMap<PathBuf, Vec<&StoredThread>> = HashMap::new();
|
||||
for entry in recent_threads {
|
||||
let group =
|
||||
resolve_root_git_project_for_trust(&entry.cwd).unwrap_or_else(|| entry.cwd.clone());
|
||||
@@ -446,7 +446,7 @@ fn format_section(title: &str, body: Option<String>, budget_tokens: usize) -> Op
|
||||
fn format_thread_group(
|
||||
current_group: &Path,
|
||||
group: &Path,
|
||||
entries: Vec<&ThreadMetadata>,
|
||||
entries: Vec<&StoredThread>,
|
||||
) -> Option<String> {
|
||||
let latest = entries.first()?;
|
||||
let group_label = if resolve_root_git_project_for_trust(latest.cwd.as_path()).is_some() {
|
||||
@@ -461,8 +461,9 @@ fn format_thread_group(
|
||||
];
|
||||
|
||||
if let Some(git_branch) = latest
|
||||
.git_branch
|
||||
.as_deref()
|
||||
.git_info
|
||||
.as_ref()
|
||||
.and_then(|git| git.branch.as_deref())
|
||||
.filter(|git_branch| !git_branch.is_empty())
|
||||
{
|
||||
lines.push(format!("Latest branch: {git_branch}"));
|
||||
|
||||
@@ -3,20 +3,31 @@ use super::build_recent_work_section;
|
||||
use super::build_workspace_section_with_user_root;
|
||||
use chrono::TimeZone;
|
||||
use chrono::Utc;
|
||||
use codex_git_utils::GitSha;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_state::ThreadMetadata;
|
||||
use codex_protocol::protocol::AskForApproval;
|
||||
use codex_protocol::protocol::GitInfo;
|
||||
use codex_protocol::protocol::SandboxPolicy;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_thread_store::StoredThread;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn thread_metadata(cwd: &str, title: &str, first_user_message: &str) -> ThreadMetadata {
|
||||
ThreadMetadata {
|
||||
id: ThreadId::new(),
|
||||
rollout_path: PathBuf::from("/tmp/rollout.jsonl"),
|
||||
fn stored_thread(cwd: &str, title: &str, first_user_message: &str) -> StoredThread {
|
||||
StoredThread {
|
||||
thread_id: ThreadId::new(),
|
||||
rollout_path: Some(PathBuf::from("/tmp/rollout.jsonl")),
|
||||
forked_from_id: None,
|
||||
preview: first_user_message.to_string(),
|
||||
name: (!title.is_empty()).then(|| title.to_string()),
|
||||
model_provider: "test-provider".to_string(),
|
||||
model: Some("gpt-5".to_string()),
|
||||
reasoning_effort: None,
|
||||
created_at: Utc
|
||||
.timestamp_opt(1_709_251_100, 0)
|
||||
.single()
|
||||
@@ -25,24 +36,23 @@ fn thread_metadata(cwd: &str, title: &str, first_user_message: &str) -> ThreadMe
|
||||
.timestamp_opt(1_709_251_200, 0)
|
||||
.single()
|
||||
.expect("valid timestamp"),
|
||||
source: "cli".to_string(),
|
||||
agent_path: None,
|
||||
agent_nickname: None,
|
||||
agent_role: None,
|
||||
model_provider: "test-provider".to_string(),
|
||||
model: Some("gpt-5".to_string()),
|
||||
reasoning_effort: None,
|
||||
archived_at: None,
|
||||
cwd: PathBuf::from(cwd),
|
||||
cli_version: "test".to_string(),
|
||||
title: title.to_string(),
|
||||
sandbox_policy: "workspace-write".to_string(),
|
||||
approval_mode: "never".to_string(),
|
||||
tokens_used: 0,
|
||||
source: SessionSource::Cli,
|
||||
agent_nickname: None,
|
||||
agent_role: None,
|
||||
agent_path: None,
|
||||
git_info: Some(GitInfo {
|
||||
commit_hash: Some(GitSha::new("abcdef")),
|
||||
branch: Some("main".to_string()),
|
||||
repository_url: None,
|
||||
}),
|
||||
approval_mode: AskForApproval::Never,
|
||||
sandbox_policy: SandboxPolicy::new_read_only_policy(),
|
||||
token_usage: None,
|
||||
first_user_message: Some(first_user_message.to_string()),
|
||||
archived_at: None,
|
||||
git_sha: None,
|
||||
git_branch: Some("main".to_string()),
|
||||
git_origin_url: None,
|
||||
history: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -224,17 +234,17 @@ fn recent_work_section_groups_threads_by_cwd() {
|
||||
fs::create_dir_all(&outside).expect("create outside dir");
|
||||
|
||||
let recent_threads = vec![
|
||||
thread_metadata(
|
||||
stored_thread(
|
||||
workspace_a.to_string_lossy().as_ref(),
|
||||
"Investigate realtime startup context",
|
||||
"Log the startup context before sending it",
|
||||
),
|
||||
thread_metadata(
|
||||
stored_thread(
|
||||
workspace_b.to_string_lossy().as_ref(),
|
||||
"Trim websocket startup payload",
|
||||
"Remove memories from the realtime startup context",
|
||||
),
|
||||
thread_metadata(outside.to_string_lossy().as_ref(), "", "Inspect flaky test"),
|
||||
stored_thread(outside.to_string_lossy().as_ref(), "", "Inspect flaky test"),
|
||||
];
|
||||
let current_cwd = workspace_a;
|
||||
let repo = fs::canonicalize(repo).expect("canonicalize repo");
|
||||
|
||||
@@ -46,11 +46,7 @@ impl codex_rollout::RolloutConfigView for Config {
|
||||
}
|
||||
|
||||
pub(crate) mod list {
|
||||
pub use codex_rollout::ThreadListConfig;
|
||||
pub use codex_rollout::ThreadListLayout;
|
||||
pub use codex_rollout::ThreadSortKey;
|
||||
pub use codex_rollout::find_thread_path_by_id_str;
|
||||
pub use codex_rollout::get_threads_in_root;
|
||||
}
|
||||
|
||||
pub(crate) mod metadata {
|
||||
|
||||
@@ -24,6 +24,7 @@ use codex_mcp::McpConnectionManager;
|
||||
use codex_models_manager::manager::ModelsManager;
|
||||
use codex_otel::SessionTelemetry;
|
||||
use codex_rollout::state_db::StateDbHandle;
|
||||
use codex_thread_store::LocalThreadStore;
|
||||
use std::path::PathBuf;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::RwLock;
|
||||
@@ -59,6 +60,7 @@ pub(crate) struct SessionServices {
|
||||
pub(crate) network_proxy: Option<StartedNetworkProxy>,
|
||||
pub(crate) network_approval: Arc<NetworkApprovalService>,
|
||||
pub(crate) state_db: Option<StateDbHandle>,
|
||||
pub(crate) thread_store: LocalThreadStore,
|
||||
/// Session-scoped model client shared across turns.
|
||||
pub(crate) model_client: ModelClient,
|
||||
pub(crate) code_mode_service: CodeModeService,
|
||||
|
||||
Reference in New Issue
Block a user