Thread state DB through rollout lookups

Remove the one-line rollout path wrappers and make callers pass the shared optional state DB handle directly. Reuse the initialized app-server handle on cold thread paths and carry existing thread DB handles through core and app-server callsites.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Charles Cunningham
2026-03-24 20:54:21 -07:00
co-authored by Codex
parent 74ec7c1be5
commit 504afb1dcd
18 changed files with 268 additions and 231 deletions
+137 -130
View File
@@ -210,16 +210,17 @@ use codex_core::exec::ExecCapturePolicy;
use codex_core::exec::ExecExpiration;
use codex_core::exec::ExecParams;
use codex_core::exec_env::create_env;
use codex_core::find_archived_thread_path_by_id_str_with_state_db;
use codex_core::find_archived_thread_path_by_id_str;
use codex_core::find_thread_name_by_id;
use codex_core::find_thread_names_by_ids;
use codex_core::find_thread_path_by_id_str_with_state_db;
use codex_core::find_thread_path_by_id_str;
use codex_core::mcp::auth::discover_supported_scopes;
use codex_core::mcp::auth::resolve_oauth_scopes;
use codex_core::mcp::collect_mcp_snapshot;
use codex_core::mcp::group_tools_by_server;
use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig;
use codex_core::parse_cursor;
use codex_core::path_utils::normalize_for_path_comparison;
use codex_core::plugins::MarketplaceError;
use codex_core::plugins::MarketplacePluginSource;
use codex_core::plugins::OPENAI_CURATED_MARKETPLACE_NAME;
@@ -234,7 +235,6 @@ use codex_core::read_session_meta_line;
use codex_core::rollout_date_parts;
use codex_core::sandboxing::SandboxPermissions;
use codex_core::state_db::StateDbHandle;
use codex_core::state_db::get_state_db;
use codex_core::state_db::reconcile_rollout;
use codex_core::windows_sandbox::WindowsSandboxLevelExt;
use codex_core::windows_sandbox::WindowsSandboxSetupMode as CoreWindowsSandboxSetupMode;
@@ -300,7 +300,6 @@ use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::SystemTime;
use tokio::sync::Mutex;
use tokio::sync::OnceCell;
use tokio::sync::broadcast;
use tokio::sync::oneshot;
use tokio::sync::watch;
@@ -385,7 +384,7 @@ pub(crate) struct CodexMessageProcessor {
pending_fuzzy_searches: Arc<Mutex<HashMap<String, Arc<AtomicBool>>>>,
fuzzy_search_sessions: Arc<Mutex<HashMap<String, FuzzyFileSearchSession>>>,
background_tasks: TaskTracker,
app_server_state_db_cell: OnceCell<StateDbHandle>,
state_db: StateDbHandle,
feedback: CodexFeedback,
log_db: Option<LogDbLayer>,
}
@@ -432,6 +431,7 @@ pub(crate) struct CodexMessageProcessorArgs {
pub(crate) cloud_requirements: Arc<RwLock<CloudRequirementsLoader>>,
pub(crate) feedback: CodexFeedback,
pub(crate) log_db: Option<LogDbLayer>,
pub(crate) state_db: StateDbHandle,
}
impl CodexMessageProcessor {
@@ -485,60 +485,6 @@ impl CodexMessageProcessor {
Ok((thread_id, thread))
}
async fn app_server_state_db(&self) -> Option<StateDbHandle> {
let config = Arc::clone(&self.config);
self.app_server_state_db_cell
.get_or_try_init(|| async move { get_state_db(&config).await.ok_or(()) })
.await
.ok()
.cloned()
}
async fn find_thread_path_by_id_with_state_db(
&self,
thread_id: ThreadId,
state_db_ctx: Option<StateDbHandle>,
) -> std::io::Result<Option<PathBuf>> {
find_thread_path_by_id_str_with_state_db(
&self.config.codex_home,
&thread_id.to_string(),
state_db_ctx,
)
.await
}
async fn find_archived_thread_path_by_id_with_state_db(
&self,
thread_id: ThreadId,
state_db_ctx: Option<StateDbHandle>,
) -> std::io::Result<Option<PathBuf>> {
find_archived_thread_path_by_id_str_with_state_db(
&self.config.codex_home,
&thread_id.to_string(),
state_db_ctx,
)
.await
}
async fn find_thread_path_by_id(
&self,
thread_id: ThreadId,
) -> std::io::Result<Option<PathBuf>> {
self.find_thread_path_by_id_with_state_db(thread_id, self.app_server_state_db().await)
.await
}
async fn find_archived_thread_path_by_id(
&self,
thread_id: ThreadId,
) -> std::io::Result<Option<PathBuf>> {
self.find_archived_thread_path_by_id_with_state_db(
thread_id,
self.app_server_state_db().await,
)
.await
}
pub fn new(args: CodexMessageProcessorArgs) -> Self {
let CodexMessageProcessorArgs {
auth_manager,
@@ -551,6 +497,7 @@ impl CodexMessageProcessor {
cloud_requirements,
feedback,
log_db,
state_db,
} = args;
Self {
auth_manager,
@@ -569,7 +516,7 @@ impl CodexMessageProcessor {
pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())),
fuzzy_search_sessions: Arc::new(Mutex::new(HashMap::new())),
background_tasks: TaskTracker::new(),
app_server_state_db_cell: OnceCell::new(),
state_db,
feedback,
log_db,
}
@@ -2294,7 +2241,13 @@ impl CodexMessageProcessor {
}
};
let rollout_path = match self.find_thread_path_by_id(thread_id).await {
let rollout_path = match find_thread_path_by_id_str(
&self.config.codex_home,
&thread_id.to_string(),
Some(self.state_db.clone()),
)
.await
{
Ok(Some(p)) => p,
Ok(None) => {
let error = JSONRPCErrorError {
@@ -2442,7 +2395,13 @@ impl CodexMessageProcessor {
return;
}
let thread_exists = match self.find_thread_path_by_id(thread_id).await {
let thread_exists = match find_thread_path_by_id_str(
&self.config.codex_home,
&thread_id.to_string(),
Some(self.state_db.clone()),
)
.await
{
Ok(Some(_)) => true,
Ok(None) => false,
Err(err) => {
@@ -2526,7 +2485,7 @@ impl CodexMessageProcessor {
let loaded_thread = self.thread_manager.get_thread(thread_uuid).await.ok();
let mut state_db_ctx = loaded_thread.as_ref().and_then(|thread| thread.state_db());
if state_db_ctx.is_none() {
state_db_ctx = self.app_server_state_db().await;
state_db_ctx = Some(self.state_db.clone());
}
let Some(state_db_ctx) = state_db_ctx else {
self.send_internal_error(
@@ -2729,17 +2688,20 @@ impl CodexMessageProcessor {
return Ok(());
}
let rollout_path = match self
.find_thread_path_by_id_with_state_db(thread_uuid, Some(state_db_ctx.clone()))
.await
let rollout_path = match find_thread_path_by_id_str(
&self.config.codex_home,
&thread_uuid.to_string(),
Some(state_db_ctx.clone()),
)
.await
{
Ok(Some(path)) => path,
Ok(None) => match self
.find_archived_thread_path_by_id_with_state_db(
thread_uuid,
Some(state_db_ctx.clone()),
)
.await
Ok(None) => match find_archived_thread_path_by_id_str(
&self.config.codex_home,
&thread_uuid.to_string(),
Some(state_db_ctx.clone()),
)
.await
{
Ok(Some(path)) => path,
Ok(None) => {
@@ -2799,7 +2761,13 @@ impl CodexMessageProcessor {
}
};
let archived_path = match self.find_archived_thread_path_by_id(thread_id).await {
let archived_path = match find_archived_thread_path_by_id_str(
&self.config.codex_home,
&thread_id.to_string(),
Some(self.state_db.clone()),
)
.await
{
Ok(Some(path)) => path,
Ok(None) => {
let error = JSONRPCErrorError {
@@ -2823,7 +2791,7 @@ impl CodexMessageProcessor {
let rollout_path_display = archived_path.display().to_string();
let fallback_provider = self.config.model_provider_id.clone();
let state_db_ctx = self.app_server_state_db().await;
let state_db_ctx = Some(self.state_db.clone());
let archived_folder = self
.config
.codex_home
@@ -3303,24 +3271,25 @@ impl CodexMessageProcessor {
let loaded_thread = self.thread_manager.get_thread(thread_uuid).await.ok();
let loaded_thread_state_db = loaded_thread.as_ref().and_then(|thread| thread.state_db());
let app_server_state_db = if loaded_thread_state_db.is_some() {
let state_db = if loaded_thread_state_db.is_some() {
None
} else {
self.app_server_state_db().await
Some(self.state_db.clone())
};
let db_summary = if let Some(state_db_ctx) = loaded_thread_state_db.as_ref() {
read_summary_from_state_db_context_by_thread_id(Some(state_db_ctx), thread_uuid).await
} else {
read_summary_from_state_db_by_thread_id(app_server_state_db.as_ref(), thread_uuid).await
read_summary_from_state_db_by_thread_id(state_db.as_ref(), thread_uuid).await
};
let mut rollout_path = db_summary.as_ref().map(|summary| summary.path.clone());
if rollout_path.is_none() || include_turns {
let state_db_ctx = loaded_thread_state_db
.clone()
.or_else(|| app_server_state_db.clone());
rollout_path = match self
.find_thread_path_by_id_with_state_db(thread_uuid, state_db_ctx)
.await
let state_db_ctx = loaded_thread_state_db.clone().or_else(|| state_db.clone());
rollout_path = match find_thread_path_by_id_str(
&self.config.codex_home,
&thread_uuid.to_string(),
state_db_ctx,
)
.await
{
Ok(Some(path)) => Some(path),
Ok(None) => {
@@ -3709,7 +3678,7 @@ impl CodexMessageProcessor {
let InitialHistory::Resumed(resumed_history) = thread_history else {
return None;
};
let state_db_ctx = self.app_server_state_db().await?;
let state_db_ctx = &self.state_db;
let persisted_metadata = state_db_ctx
.get_thread(resumed_history.conversation_id)
.await
@@ -3742,7 +3711,15 @@ impl CodexMessageProcessor {
if path.exists() {
path
} else {
match self.find_thread_path_by_id(existing_thread_id).await {
match find_thread_path_by_id_str(
&self.config.codex_home,
&existing_thread_id.to_string(),
existing_thread
.state_db()
.or_else(|| Some(self.state_db.clone())),
)
.await
{
Ok(Some(path)) => path,
Ok(None) => {
self.send_invalid_request_error(
@@ -3763,7 +3740,15 @@ impl CodexMessageProcessor {
}
}
} else {
match self.find_thread_path_by_id(existing_thread_id).await {
match find_thread_path_by_id_str(
&self.config.codex_home,
&existing_thread_id.to_string(),
existing_thread
.state_db()
.or_else(|| Some(self.state_db.clone())),
)
.await
{
Ok(Some(path)) => path,
Ok(None) => {
self.send_invalid_request_error(
@@ -3917,7 +3902,13 @@ impl CodexMessageProcessor {
}
};
match self.find_thread_path_by_id(existing_thread_id).await {
match find_thread_path_by_id_str(
&self.config.codex_home,
&existing_thread_id.to_string(),
Some(self.state_db.clone()),
)
.await
{
Ok(Some(path)) => path,
Ok(None) => {
self.send_invalid_request_error(
@@ -4044,7 +4035,13 @@ impl CodexMessageProcessor {
}
};
match self.find_thread_path_by_id(existing_thread_id).await {
match find_thread_path_by_id_str(
&self.config.codex_home,
&existing_thread_id.to_string(),
Some(self.state_db.clone()),
)
.await
{
Ok(Some(p)) => (p, Some(existing_thread_id)),
Ok(None) => {
self.send_invalid_request_error(
@@ -4065,9 +4062,9 @@ impl CodexMessageProcessor {
}
};
let app_server_state_db = self.app_server_state_db().await;
let state_db = &self.state_db;
let history_cwd = read_history_cwd_from_state_db(
app_server_state_db.as_ref(),
Some(state_db),
source_thread_id,
rollout_path.as_path(),
)
@@ -4293,13 +4290,10 @@ impl CodexMessageProcessor {
request_id: ConnectionRequestId,
params: GetConversationSummaryParams,
) {
let app_server_state_db = self.app_server_state_db().await;
let state_db = &self.state_db;
if let GetConversationSummaryParams::ThreadId { conversation_id } = &params
&& let Some(summary) = read_summary_from_state_db_by_thread_id(
app_server_state_db.as_ref(),
*conversation_id,
)
.await
&& let Some(summary) =
read_summary_from_state_db_by_thread_id(Some(state_db), *conversation_id).await
{
let response = GetConversationSummaryResponse { summary };
self.outgoing.send_response(request_id, response).await;
@@ -4315,12 +4309,12 @@ impl CodexMessageProcessor {
}
}
GetConversationSummaryParams::ThreadId { conversation_id } => {
match self
.find_thread_path_by_id_with_state_db(
conversation_id,
app_server_state_db.clone(),
)
.await
match find_thread_path_by_id_str(
&self.config.codex_home,
&conversation_id.to_string(),
Some(self.state_db.clone()),
)
.await
{
Ok(Some(p)) => p,
_ => {
@@ -4401,7 +4395,11 @@ impl CodexMessageProcessor {
let fallback_provider = self.config.model_provider_id.clone();
let (allowed_sources_vec, source_kind_filter) = compute_source_filters(source_kinds);
let allowed_sources = allowed_sources_vec.as_slice();
let state_db_ctx = self.app_server_state_db().await;
let state_db_ctx = Some(self.state_db.clone());
let normalized_cwd = cwd
.as_ref()
.and_then(|cwd| normalize_for_path_comparison(cwd).ok());
let cwd_filter = normalized_cwd.as_ref().or(cwd.as_ref());
while remaining > 0 {
let page_size = remaining.min(THREAD_LIST_MAX_LIMIT);
@@ -4443,7 +4441,7 @@ impl CodexMessageProcessor {
let mut filtered = Vec::with_capacity(page.items.len());
for it in page.items {
let Some(summary) = summary_from_thread_list_item(
let Some(mut summary) = summary_from_thread_list_item(
it,
fallback_provider.as_str(),
state_db_ctx.as_ref(),
@@ -4455,10 +4453,11 @@ impl CodexMessageProcessor {
if source_kind_filter
.as_ref()
.is_none_or(|filter| source_kind_matches(&summary.source, filter))
&& cwd
.as_ref()
.is_none_or(|expected_cwd| &summary.cwd == expected_cwd)
&& cwd_filter.is_none_or(|expected_cwd| &summary.cwd == expected_cwd)
{
if let Some(cwd) = cwd.as_ref() {
summary.cwd = cwd.clone();
}
filtered.push(summary);
if filtered.len() >= remaining {
break;
@@ -5254,7 +5253,7 @@ impl CodexMessageProcessor {
self.finalize_thread_teardown(thread_id).await;
if state_db_ctx.is_none() {
state_db_ctx = self.app_server_state_db().await;
state_db_ctx = Some(self.state_db.clone());
}
// Move the rollout file to archived.
@@ -6605,18 +6604,24 @@ impl CodexMessageProcessor {
let rollout_path = if let Some(path) = parent_thread.rollout_path() {
path
} else {
self.find_thread_path_by_id(parent_thread_id)
.await
.map_err(|err| JSONRPCErrorError {
code: INTERNAL_ERROR_CODE,
message: format!("failed to locate thread id {parent_thread_id}: {err}"),
data: None,
})?
.ok_or_else(|| JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("no rollout found for thread id {parent_thread_id}"),
data: None,
})?
find_thread_path_by_id_str(
&self.config.codex_home,
&parent_thread_id.to_string(),
parent_thread
.state_db()
.or_else(|| Some(self.state_db.clone())),
)
.await
.map_err(|err| JSONRPCErrorError {
code: INTERNAL_ERROR_CODE,
message: format!("failed to locate thread id {parent_thread_id}: {err}"),
data: None,
})?
.ok_or_else(|| JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("no rollout found for thread id {parent_thread_id}"),
data: None,
})?
};
let mut config = self.config.as_ref().clone();
@@ -7209,11 +7214,10 @@ impl CodexMessageProcessor {
if let Some(log_db) = self.log_db.as_ref() {
log_db.flush().await;
}
let state_db_ctx = self.app_server_state_db().await;
match (state_db_ctx.as_ref(), conversation_id) {
(Some(state_db_ctx), Some(conversation_id)) => {
match conversation_id {
Some(conversation_id) => {
let thread_id_text = conversation_id.to_string();
match state_db_ctx.query_feedback_logs(&thread_id_text).await {
match self.state_db.query_feedback_logs(&thread_id_text).await {
Ok(logs) if logs.is_empty() => None,
Ok(logs) => Some(logs),
Err(err) => {
@@ -7224,7 +7228,7 @@ impl CodexMessageProcessor {
}
}
}
_ => None,
None => None,
}
} else {
None
@@ -7362,11 +7366,14 @@ impl CodexMessageProcessor {
async fn resolve_rollout_path(&self, conversation_id: ThreadId) -> Option<PathBuf> {
match self.thread_manager.get_thread(conversation_id).await {
Ok(conv) => conv.rollout_path(),
Err(_) => self
.find_thread_path_by_id(conversation_id)
.await
.ok()
.flatten(),
Err(_) => find_thread_path_by_id_str(
&self.config.codex_home,
&conversation_id.to_string(),
Some(self.state_db.clone()),
)
.await
.ok()
.flatten(),
}
}
}
+20 -7
View File
@@ -77,6 +77,8 @@ use codex_arg0::Arg0DispatchPaths;
use codex_core::config::Config;
use codex_core::config_loader::CloudRequirementsLoader;
use codex_core::config_loader::LoaderOverrides;
use codex_core::state_db::StateDbHandle;
use codex_core::state_db::init;
use codex_feedback::CodexFeedback;
use codex_protocol::protocol::SessionSource;
use tokio::sync::mpsc;
@@ -325,7 +327,13 @@ impl InProcessClientHandle {
/// the runtime is shut down and an `InvalidData` error is returned.
pub async fn start(args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
let initialize = args.initialize.clone();
let client = start_uninitialized(args);
let state_db = init(args.config.as_ref()).await.ok_or_else(|| {
IoError::other(format!(
"failed to initialize state db at {}",
args.config.sqlite_home.display()
))
})?;
let client = start_uninitialized(args, state_db);
let initialize_response = client
.request(ClientRequest::Initialize {
@@ -345,7 +353,7 @@ pub async fn start(args: InProcessStartArgs) -> IoResult<InProcessClientHandle>
Ok(client)
}
fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
fn start_uninitialized(args: InProcessStartArgs, state_db: StateDbHandle) -> InProcessClientHandle {
let channel_capacity = args.channel_capacity.max(1);
let (client_tx, mut client_rx) = mpsc::channel::<InProcessClientMessage>(channel_capacity);
let (event_tx, event_rx) = mpsc::channel::<InProcessServerEvent>(channel_capacity);
@@ -388,6 +396,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
cloud_requirements: args.cloud_requirements,
feedback: args.feedback,
log_db: None,
state_db,
config_warnings: args.config_warnings,
session_source: args.session_source,
enable_codex_api_key_env: args.enable_codex_api_key_env,
@@ -693,13 +702,17 @@ mod tests {
use codex_app_server_protocol::TurnStatus;
use codex_core::config::ConfigBuilder;
use pretty_assertions::assert_eq;
use tempfile::TempDir;
async fn build_test_config() -> Config {
match ConfigBuilder::default().build().await {
Ok(config) => config,
Err(_) => Config::load_default_with_cli_overrides(Vec::new())
.expect("default config should load"),
}
let codex_home = TempDir::new().expect("create temp dir");
let config = ConfigBuilder::default()
.codex_home(codex_home.path().to_path_buf())
.build()
.await
.expect("test config should build");
std::mem::forget(codex_home);
config
}
async fn start_test_client_with_capacity(
+10 -7
View File
@@ -8,9 +8,11 @@ use codex_core::config::ConfigBuilder;
use codex_core::config_loader::CloudRequirementsLoader;
use codex_core::config_loader::ConfigLayerStackOrdering;
use codex_core::config_loader::LoaderOverrides;
use codex_core::state_db::init;
use codex_utils_cli::CliConfigOverrides;
use std::collections::HashMap;
use std::collections::HashSet;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::io::Result as IoResult;
use std::sync::Arc;
@@ -489,6 +491,12 @@ pub async fn run_main_with_transport(
}
let feedback = CodexFeedback::new();
let state_db = init(&config).await.ok_or_else(|| {
IoError::other(format!(
"failed to initialize state db at {}",
config.sqlite_home.display()
))
})?;
let otel = codex_core::otel_init::build_provider(
&config,
@@ -522,13 +530,7 @@ pub async fn run_main_with_transport(
let feedback_layer = feedback.logger_layer();
let feedback_metadata_layer = feedback.metadata_layer();
let log_db = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
)
.await
.ok()
.map(log_db::start);
let log_db = Some(log_db::start(Arc::clone(&state_db)));
let log_db_layer = log_db
.clone()
.map(|layer| layer.with_filter(Targets::new().with_default(Level::TRACE)));
@@ -618,6 +620,7 @@ pub async fn run_main_with_transport(
cloud_requirements: cloud_requirements.clone(),
feedback: feedback.clone(),
log_db,
state_db,
config_warnings,
session_source,
enable_codex_api_key_env: false,
@@ -66,6 +66,7 @@ use codex_core::default_client::get_codex_user_agent;
use codex_core::default_client::set_default_client_residency_requirement;
use codex_core::default_client::set_default_originator;
use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig;
use codex_core::state_db::StateDbHandle;
use codex_features::Feature;
use codex_feedback::CodexFeedback;
use codex_login::auth::ExternalAuthRefreshContext;
@@ -181,6 +182,7 @@ pub(crate) struct MessageProcessorArgs {
pub(crate) cloud_requirements: CloudRequirementsLoader,
pub(crate) feedback: CodexFeedback,
pub(crate) log_db: Option<LogDbLayer>,
pub(crate) state_db: StateDbHandle,
pub(crate) config_warnings: Vec<ConfigWarningNotification>,
pub(crate) session_source: SessionSource,
pub(crate) enable_codex_api_key_env: bool,
@@ -199,6 +201,7 @@ impl MessageProcessor {
cloud_requirements,
feedback,
log_db,
state_db,
config_warnings,
session_source,
enable_codex_api_key_env,
@@ -242,6 +245,7 @@ impl MessageProcessor {
cloud_requirements: cloud_requirements.clone(),
feedback,
log_db,
state_db,
});
// Keep plugin startup warmups aligned at app-server startup.
// TODO(xl): Move into PluginManager once this no longer depends on config feature gating.
@@ -24,6 +24,7 @@ use codex_core::config::Config;
use codex_core::config::ConfigBuilder;
use codex_core::config_loader::CloudRequirementsLoader;
use codex_core::config_loader::LoaderOverrides;
use codex_core::state_db::init;
use codex_feedback::CodexFeedback;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::W3cTraceContext;
@@ -117,7 +118,7 @@ impl TracingHarness {
let server = create_mock_responses_server_repeating_assistant("Done").await;
let codex_home = TempDir::new()?;
let config = Arc::new(build_test_config(codex_home.path(), &server.uri()).await?);
let (processor, outgoing_rx) = build_test_processor(config);
let (processor, outgoing_rx) = build_test_processor(config).await;
let tracing = init_test_tracing();
tracing.exporter.reset();
tracing::callsite::rebuild_interest_cache();
@@ -224,7 +225,7 @@ async fn build_test_config(codex_home: &Path, server_uri: &str) -> Result<Config
.await?)
}
fn build_test_processor(
async fn build_test_processor(
config: Arc<Config>,
) -> (
MessageProcessor,
@@ -232,6 +233,9 @@ fn build_test_processor(
) {
let (outgoing_tx, outgoing_rx) = mpsc::channel(16);
let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx));
let state_db = init(config.as_ref())
.await
.expect("state db should initialize in tracing tests");
let processor = MessageProcessor::new(MessageProcessorArgs {
outgoing,
arg0_paths: Arg0DispatchPaths::default(),
@@ -241,6 +245,7 @@ fn build_test_processor(
cloud_requirements: CloudRequirementsLoader::default(),
feedback: CodexFeedback::new(),
log_db: None,
state_db,
config_warnings: Vec::new(),
session_source: SessionSource::VSCode,
enable_codex_api_key_env: false,
@@ -58,7 +58,7 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> {
rollout_path.display()
);
assert!(
find_thread_path_by_id_str(codex_home.path(), &thread.id)
find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None)
.await?
.is_none(),
"thread id should not be discoverable before rollout materialization"
@@ -113,9 +113,10 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> {
rollout_path.display()
);
let discovered_path = find_thread_path_by_id_str(codex_home.path(), &thread.id)
.await?
.expect("expected rollout path for thread id to exist after materialization");
let discovered_path =
find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None)
.await?
.expect("expected rollout path for thread id to exist after materialization");
assert_paths_match_on_disk(&discovered_path, &rollout_path)?;
let archive_id = mcp
@@ -75,9 +75,10 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result
)
.await??;
let found_rollout_path = find_thread_path_by_id_str(codex_home.path(), &thread.id)
.await?
.expect("expected rollout path for thread id to exist");
let found_rollout_path =
find_thread_path_by_id_str(codex_home.path(), &thread.id, /*state_db_ctx*/ None)
.await?
.expect("expected rollout path for thread id to exist");
assert_paths_match_on_disk(&found_rollout_path, &rollout_path)?;
let archive_id = mcp
@@ -92,9 +93,13 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result
.await??;
let _: ThreadArchiveResponse = to_response::<ThreadArchiveResponse>(archive_resp)?;
let archived_path = find_archived_thread_path_by_id_str(codex_home.path(), &thread.id)
.await?
.expect("expected archived rollout path for thread id to exist");
let archived_path = find_archived_thread_path_by_id_str(
codex_home.path(),
&thread.id,
/*state_db_ctx*/ None,
)
.await?
.expect("expected archived rollout path for thread id to exist");
let archived_path_display = archived_path.display();
assert!(
archived_path.exists(),
+21 -13
View File
@@ -205,6 +205,9 @@ impl AgentControl {
.or(find_thread_path_by_id_str(
config.codex_home.as_path(),
&parent_thread_id.to_string(),
parent_thread
.as_ref()
.and_then(|parent_thread| parent_thread.state_db()),
)
.await?)
.ok_or_else(|| {
@@ -377,6 +380,7 @@ impl AgentControl {
}
let state = self.upgrade()?;
let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?;
let state_db_ctx = state_db::get_state_db(&config).await;
let (session_source, agent_metadata) = match session_source {
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
@@ -386,7 +390,7 @@ impl AgentControl {
agent_nickname: _,
}) => {
let (resumed_agent_nickname, resumed_agent_role) =
if let Some(state_db_ctx) = state_db::get_state_db(&config).await {
if let Some(state_db_ctx) = state_db_ctx.as_ref() {
match state_db_ctx.get_thread(thread_id).await {
Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role),
Ok(None) | Err(_) => (None, None),
@@ -413,18 +417,22 @@ impl AgentControl {
let inherited_exec_policy = self
.inherited_exec_policy_for_source(&state, Some(&session_source), &config)
.await;
let rollout_path =
match find_thread_path_by_id_str(config.codex_home.as_path(), &thread_id.to_string())
.await?
{
Some(rollout_path) => rollout_path,
None => find_archived_thread_path_by_id_str(
config.codex_home.as_path(),
&thread_id.to_string(),
)
.await?
.ok_or_else(|| CodexErr::ThreadNotFound(thread_id))?,
};
let rollout_path = match find_thread_path_by_id_str(
config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.clone(),
)
.await?
{
Some(rollout_path) => rollout_path,
None => find_archived_thread_path_by_id_str(
config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx,
)
.await?
.ok_or_else(|| CodexErr::ThreadNotFound(thread_id))?,
};
let resumed_thread = state
.resume_thread_from_rollout_with_source(
-2
View File
@@ -144,12 +144,10 @@ pub use rollout::SESSIONS_SUBDIR;
pub use rollout::SessionMeta;
pub use rollout::append_thread_name;
pub use rollout::find_archived_thread_path_by_id_str;
pub use rollout::find_archived_thread_path_by_id_str_with_state_db;
#[deprecated(note = "use find_thread_path_by_id_str")]
pub use rollout::find_conversation_path_by_id_str;
pub use rollout::find_thread_name_by_id;
pub use rollout::find_thread_path_by_id_str;
pub use rollout::find_thread_path_by_id_str_with_state_db;
pub use rollout::find_thread_path_by_name_str;
pub use rollout::list::Cursor;
pub use rollout::list::ThreadItem;
+2 -1
View File
@@ -518,7 +518,8 @@ pub async fn cleanup_stale_snapshots(codex_home: &Path, active_session_id: Threa
continue;
}
let rollout_path = find_thread_path_by_id_str(codex_home, session_id).await?;
let rollout_path =
find_thread_path_by_id_str(codex_home, session_id, /*state_db_ctx*/ None).await?;
let Some(rollout_path) = rollout_path else {
remove_snapshot_file(&path).await;
continue;
+27 -18
View File
@@ -78,9 +78,10 @@ async fn find_locates_rollout_file_by_id() {
let id = Uuid::new_v4();
let expected = write_minimal_rollout_with_id(home.path(), id);
let found = find_thread_path_by_id_str(home.path(), &id.to_string())
.await
.unwrap();
let found =
find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None)
.await
.unwrap();
assert_eq!(found.unwrap(), expected);
}
@@ -94,9 +95,10 @@ async fn find_handles_gitignore_covering_codex_home_directory() {
let id = Uuid::new_v4();
let expected = write_minimal_rollout_with_id(&codex_home, id);
let found = find_thread_path_by_id_str(&codex_home, &id.to_string())
.await
.unwrap();
let found =
find_thread_path_by_id_str(&codex_home, &id.to_string(), /*state_db_ctx*/ None)
.await
.unwrap();
assert_eq!(found, Some(expected));
}
@@ -114,9 +116,10 @@ async fn find_prefers_sqlite_path_by_id() {
write_minimal_rollout_with_id(home.path(), id);
upsert_thread_metadata(home.path(), thread_id, db_path.clone()).await;
let found = find_thread_path_by_id_str(home.path(), &id.to_string())
.await
.unwrap();
let found =
find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None)
.await
.unwrap();
assert_eq!(found, Some(db_path));
}
@@ -133,9 +136,10 @@ async fn find_falls_back_to_filesystem_when_sqlite_has_no_match() {
.join("sessions/2030/12/30/rollout-2030-12-30T00-00-00-unrelated.jsonl");
upsert_thread_metadata(home.path(), unrelated_thread_id, unrelated_path).await;
let found = find_thread_path_by_id_str(home.path(), &id.to_string())
.await
.unwrap();
let found =
find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None)
.await
.unwrap();
assert_eq!(found, Some(expected));
}
@@ -147,9 +151,10 @@ async fn find_ignores_granular_gitignore_rules() {
let expected = write_minimal_rollout_with_id(home.path(), id);
std::fs::write(home.path().join("sessions/.gitignore"), "*.jsonl\n").unwrap();
let found = find_thread_path_by_id_str(home.path(), &id.to_string())
.await
.unwrap();
let found =
find_thread_path_by_id_str(home.path(), &id.to_string(), /*state_db_ctx*/ None)
.await
.unwrap();
assert_eq!(found, Some(expected));
}
@@ -210,9 +215,13 @@ async fn find_archived_locates_rollout_file_by_id() {
let id = Uuid::new_v4();
let expected = write_minimal_rollout_with_id_in_subdir(home.path(), "archived_sessions", id);
let found = find_archived_thread_path_by_id_str(home.path(), &id.to_string())
.await
.unwrap();
let found = find_archived_thread_path_by_id_str(
home.path(),
&id.to_string(),
/*state_db_ctx*/ None,
)
.await
.unwrap();
assert_eq!(found, Some(expected));
}
-2
View File
@@ -33,11 +33,9 @@ pub use codex_protocol::protocol::SessionMeta;
pub use config::RolloutConfig;
pub use config::RolloutConfigView;
pub use list::find_archived_thread_path_by_id_str;
pub use list::find_archived_thread_path_by_id_str_with_state_db;
pub use list::find_thread_path_by_id_str;
#[deprecated(note = "use find_thread_path_by_id_str")]
pub use list::find_thread_path_by_id_str as find_conversation_path_by_id_str;
pub use list::find_thread_path_by_id_str_with_state_db;
pub use list::rollout_date_parts;
pub use policy::EventPersistenceMode;
pub use recorder::RolloutRecorder;
+4 -34
View File
@@ -1251,51 +1251,21 @@ async fn find_thread_path_by_id_str_in_subdir(
}
/// Locate a recorded thread rollout file by its UUID string using the existing
/// paginated listing implementation. Returns `Ok(Some(path))` if found, `Ok(None)` if not present
/// or the id is invalid.
/// paginated listing implementation. Uses `state_db_ctx` when available to avoid reopening
/// SQLite. Returns `Ok(Some(path))` if found, `Ok(None)` if not present or the id is invalid.
pub async fn find_thread_path_by_id_str(
codex_home: &Path,
id_str: &str,
) -> io::Result<Option<PathBuf>> {
find_thread_path_by_id_str_in_subdir(
codex_home,
SESSIONS_SUBDIR,
id_str,
/*state_db_ctx*/ None,
)
.await
}
/// Locate a recorded thread rollout file by its UUID string using an already-open SQLite handle
/// when available. Returns `Ok(Some(path))` if found, `Ok(None)` if not present or the id is
/// invalid.
pub async fn find_thread_path_by_id_str_with_state_db(
codex_home: &Path,
id_str: &str,
state_db_ctx: Option<crate::state_db::StateDbHandle>,
) -> io::Result<Option<PathBuf>> {
find_thread_path_by_id_str_in_subdir(codex_home, SESSIONS_SUBDIR, id_str, state_db_ctx).await
}
/// Locate an archived thread rollout file by its UUID string.
/// Locate an archived thread rollout file by its UUID string. Uses `state_db_ctx` when available
/// to avoid reopening SQLite.
pub async fn find_archived_thread_path_by_id_str(
codex_home: &Path,
id_str: &str,
) -> io::Result<Option<PathBuf>> {
find_thread_path_by_id_str_in_subdir(
codex_home,
ARCHIVED_SESSIONS_SUBDIR,
id_str,
/*state_db_ctx*/ None,
)
.await
}
/// Locate an archived thread rollout file by its UUID string using an already-open SQLite handle
/// when available.
pub async fn find_archived_thread_path_by_id_str_with_state_db(
codex_home: &Path,
id_str: &str,
state_db_ctx: Option<crate::state_db::StateDbHandle>,
) -> io::Result<Option<PathBuf>> {
find_thread_path_by_id_str_in_subdir(codex_home, ARCHIVED_SESSIONS_SUBDIR, id_str, state_db_ctx)
+6
View File
@@ -287,6 +287,12 @@ impl RolloutRecorder {
)
.await
{
if db_page.items.is_empty() && !fs_page.items.is_empty() {
tracing::warn!(
"state db discrepancy during list_threads_with_db_fallback: empty_db_page"
);
return Ok(truncate_fs_page(fs_page, page_size, sort_key));
}
return Ok(db_page.into());
}
// If SQLite listing still fails, return the filesystem page rather than failing the list.
+6 -1
View File
@@ -139,7 +139,12 @@ pub async fn find_thread_path_by_name_str(
let Some(thread_id) = find_thread_id_by_name(codex_home, name).await? else {
return Ok(None);
};
super::list::find_thread_path_by_id_str(codex_home, &thread_id.to_string()).await
super::list::find_thread_path_by_id_str(
codex_home,
&thread_id.to_string(),
/*state_db_ctx*/ None,
)
.await
}
fn session_index_path(codex_home: &Path) -> PathBuf {
+3
View File
@@ -25,6 +25,9 @@ use uuid::Uuid;
pub type StateDbHandle = Arc<codex_state::StateRuntime>;
/// Initialize the state runtime for thread state persistence and backfill checks.
///
/// Callers that keep a shared runtime alive for repeated state reads should use
/// this helper instead of reopening SQLite on each lookup.
pub async fn init(config: &impl RolloutConfigView) -> Option<StateDbHandle> {
let config = RolloutConfig::from_view(config);
let runtime = match codex_state::StateRuntime::init(
+2 -2
View File
@@ -231,7 +231,7 @@ async fn find_thread_path_falls_back_when_db_path_is_stale() {
));
insert_state_db_thread(home, thread_id, stale_db_path.as_path(), false).await;
let found = find_thread_path_by_id_str(home, &uuid.to_string())
let found = find_thread_path_by_id_str(home, &uuid.to_string(), /*state_db_ctx*/ None)
.await
.expect("lookup should succeed");
assert_eq!(found, Some(fs_rollout_path.clone()));
@@ -257,7 +257,7 @@ async fn find_thread_path_repairs_missing_db_row_after_filesystem_fallback() {
.await
.expect("backfill should be complete");
let found = find_thread_path_by_id_str(home, &uuid.to_string())
let found = find_thread_path_by_id_str(home, &uuid.to_string(), /*state_db_ctx*/ None)
.await
.expect("lookup should succeed");
assert_eq!(found, Some(fs_rollout_path.clone()));
+3 -2
View File
@@ -721,7 +721,8 @@ async fn run_ratatui_app(
if let Some(id_str) = cli.fork_session_id.as_deref() {
let is_uuid = Uuid::parse_str(id_str).is_ok();
let path = if is_uuid {
find_thread_path_by_id_str(&config.codex_home, id_str).await?
find_thread_path_by_id_str(&config.codex_home, id_str, /*state_db_ctx*/ None)
.await?
} else {
find_thread_path_by_name_str(&config.codex_home, id_str).await?
};
@@ -814,7 +815,7 @@ async fn run_ratatui_app(
} else if let Some(id_str) = cli.resume_session_id.as_deref() {
let is_uuid = Uuid::parse_str(id_str).is_ok();
let path = if is_uuid {
find_thread_path_by_id_str(&config.codex_home, id_str).await?
find_thread_path_by_id_str(&config.codex_home, id_str, /*state_db_ctx*/ None).await?
} else {
find_thread_path_by_name_str(&config.codex_home, id_str).await?
};