chore: improve DB flushing (#13620)

This branch:
* Avoid flushing DB when not necessary
* Filter events for which we perfom an `upsert` into the DB
* Add a dedicated update function of the `thread:updated_at` that is
lighter

This should significantly reduce the DB lock contention. If it is not
sufficient, we can de-sync the flush of the DB for `updated_at`

codex-rs/core/src/codex.rs

@@ -1694,7 +1694,7 @@ impl Session {
         self.services.state_db.clone()
     }
 
-    /// Ensure all rollout writes are durably flushed.
+    /// Ensure rollout file writes are durably flushed.
     pub(crate) async fn flush_rollout(&self) {
         let recorder = {
             let guard = self.services.rollout.lock().await;

codex-rs/core/src/rollout/recorder.rs

@@ -7,6 +7,7 @@ use std::path::Path;
 use std::path::PathBuf;
 
 use chrono::SecondsFormat;
+use chrono::Utc;
 use codex_protocol::ThreadId;
 use codex_protocol::dynamic_tools::DynamicToolSpec;
 use codex_protocol::models::BaseInstructions;
@@ -448,7 +449,6 @@ impl RolloutRecorder {
         // future will yield, which is fine – we only need to ensure we do not
         // perform *blocking* I/O on the caller's thread.
         let (tx, rx) = mpsc::channel::<RolloutCmd>(256);
-
         // Spawn a Tokio task that owns the file handle and performs async
         // writes. Using `tokio::fs::File` keeps everything on the async I/O
         // driver instead of blocking the runtime.
@@ -614,14 +614,15 @@ impl RolloutRecorder {
         match self.tx.send(RolloutCmd::Shutdown { ack: tx_done }).await {
             Ok(_) => rx_done
                 .await
-                .map_err(|e| IoError::other(format!("failed waiting for rollout shutdown: {e}"))),
+                .map_err(|e| IoError::other(format!("failed waiting for rollout shutdown: {e}")))?,
             Err(e) => {
                 warn!("failed to send rollout shutdown command: {e}");
-                Err(IoError::other(format!(
+                return Err(IoError::other(format!(
                     "failed to send rollout shutdown command: {e}"
-                )))
+                )));
             }
-        }
+        };
+        Ok(())
     }
 }
 
@@ -744,25 +745,21 @@ async fn rollout_writer(
     while let Some(cmd) = rx.recv().await {
         match cmd {
             RolloutCmd::AddItems(items) => {
-                let mut persisted_items = Vec::new();
-                for item in items {
-                    persisted_items.push(item);
-                }
-                if persisted_items.is_empty() {
+                if items.is_empty() {
                     continue;
                 }
 
                 if writer.is_none() {
-                    buffered_items.extend(persisted_items);
+                    buffered_items.extend(items);
                     continue;
                 }
 
                 write_and_reconcile_items(
                     writer.as_mut(),
-                    persisted_items.as_slice(),
+                    items.as_slice(),
                     &rollout_path,
                     state_db_ctx.as_deref(),
-                    &mut state_builder,
+                    state_builder.as_ref(),
                     default_provider.as_str(),
                 )
                 .await?;
@@ -800,7 +797,7 @@ async fn rollout_writer(
                                 buffered_items.as_slice(),
                                 &rollout_path,
                                 state_db_ctx.as_deref(),
-                                &mut state_builder,
+                                state_builder.as_ref(),
                                 default_provider.as_str(),
                             )
                             .await?;
@@ -861,13 +858,12 @@ async fn write_session_meta(
     if let Some(writer) = writer.as_mut() {
         writer.write_rollout_item(&rollout_item).await?;
     }
-    state_db::reconcile_rollout(
+    sync_thread_state_after_write(
         state_db_ctx,
         rollout_path,
-        default_provider,
         state_builder.as_ref(),
         std::slice::from_ref(&rollout_item),
-        None,
+        default_provider,
         (!generate_memories).then_some("disabled"),
     )
     .await;
@@ -879,7 +875,7 @@ async fn write_and_reconcile_items(
     items: &[RolloutItem],
     rollout_path: &Path,
     state_db_ctx: Option<&StateRuntime>,
-    state_builder: &mut Option<ThreadMetadataBuilder>,
+    state_builder: Option<&ThreadMetadataBuilder>,
     default_provider: &str,
 ) -> std::io::Result<()> {
     if let Some(writer) = writer.as_mut() {
@@ -887,20 +883,65 @@ async fn write_and_reconcile_items(
             writer.write_rollout_item(item).await?;
         }
     }
-    if let Some(builder) = state_builder.as_mut() {
-        builder.rollout_path = rollout_path.to_path_buf();
+    sync_thread_state_after_write(
+        state_db_ctx,
+        rollout_path,
+        state_builder,
+        items,
+        default_provider,
+        None,
+    )
+    .await;
+    Ok(())
+}
+
+async fn sync_thread_state_after_write(
+    state_db_ctx: Option<&StateRuntime>,
+    rollout_path: &Path,
+    state_builder: Option<&ThreadMetadataBuilder>,
+    items: &[RolloutItem],
+    default_provider: &str,
+    new_thread_memory_mode: Option<&str>,
+) {
+    let updated_at = Utc::now();
+    if new_thread_memory_mode.is_some()
+        || items
+            .iter()
+            .any(codex_state::rollout_item_affects_thread_metadata)
+    {
+        state_db::apply_rollout_items(
+            state_db_ctx,
+            rollout_path,
+            default_provider,
+            state_builder,
+            items,
+            "rollout_writer",
+            new_thread_memory_mode,
+            Some(updated_at),
+        )
+        .await;
+        return;
+    }
+
+    let thread_id = state_builder
+        .map(|builder| builder.id)
+        .or_else(|| metadata::builder_from_items(items, rollout_path).map(|builder| builder.id));
+    if state_db::touch_thread_updated_at(state_db_ctx, thread_id, updated_at, "rollout_writer")
+        .await
+    {
+        return;
     }
     state_db::apply_rollout_items(
         state_db_ctx,
         rollout_path,
         default_provider,
-        state_builder.as_ref(),
+        state_builder,
         items,
         "rollout_writer",
-        None,
+        new_thread_memory_mode,
+        Some(updated_at),
     )
     .await;
-    Ok(())
 }
 
 struct JsonlWriter {
@@ -1079,6 +1120,7 @@ mod tests {
     use std::io::Write;
     use std::path::Path;
     use std::path::PathBuf;
+    use std::time::Duration;
     use tempfile::TempDir;
     use uuid::Uuid;
 
@@ -1200,6 +1242,151 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn metadata_irrelevant_events_touch_state_db_updated_at() -> std::io::Result<()> {
+        let home = TempDir::new().expect("temp dir");
+        let mut config = ConfigBuilder::default()
+            .codex_home(home.path().to_path_buf())
+            .build()
+            .await?;
+        config
+            .features
+            .enable(Feature::Sqlite)
+            .expect("test config should allow sqlite");
+
+        let state_db = codex_state::StateRuntime::init(
+            home.path().to_path_buf(),
+            config.model_provider_id.clone(),
+            None,
+        )
+        .await
+        .expect("state db should initialize");
+        state_db
+            .mark_backfill_complete(None)
+            .await
+            .expect("backfill should be complete");
+
+        let thread_id = ThreadId::new();
+        let recorder = RolloutRecorder::new(
+            &config,
+            RolloutRecorderParams::new(
+                thread_id,
+                None,
+                SessionSource::Cli,
+                BaseInstructions::default(),
+                Vec::new(),
+                EventPersistenceMode::Limited,
+            ),
+            Some(state_db.clone()),
+            None,
+        )
+        .await?;
+
+        recorder
+            .record_items(&[RolloutItem::EventMsg(EventMsg::UserMessage(
+                UserMessageEvent {
+                    message: "first-user-message".to_string(),
+                    images: None,
+                    local_images: Vec::new(),
+                    text_elements: Vec::new(),
+                },
+            ))])
+            .await?;
+        recorder.persist().await?;
+        recorder.flush().await?;
+        let initial_thread = state_db
+            .get_thread(thread_id)
+            .await
+            .expect("thread should load")
+            .expect("thread should exist");
+        let initial_updated_at = initial_thread.updated_at;
+        let initial_title = initial_thread.title.clone();
+        let initial_first_user_message = initial_thread.first_user_message.clone();
+
+        tokio::time::sleep(Duration::from_secs(1)).await;
+
+        recorder
+            .record_items(&[RolloutItem::EventMsg(EventMsg::AgentMessage(
+                AgentMessageEvent {
+                    message: "assistant text".to_string(),
+                    phase: None,
+                },
+            ))])
+            .await?;
+        recorder.flush().await?;
+
+        let updated_thread = state_db
+            .get_thread(thread_id)
+            .await
+            .expect("thread should load after agent message")
+            .expect("thread should still exist");
+
+        assert!(updated_thread.updated_at > initial_updated_at);
+        assert_eq!(updated_thread.title, initial_title);
+        assert_eq!(
+            updated_thread.first_user_message,
+            initial_first_user_message
+        );
+
+        recorder.shutdown().await?;
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn metadata_irrelevant_events_fall_back_to_upsert_when_thread_missing()
+    -> std::io::Result<()> {
+        let home = TempDir::new().expect("temp dir");
+        let mut config = ConfigBuilder::default()
+            .codex_home(home.path().to_path_buf())
+            .build()
+            .await?;
+        config
+            .features
+            .enable(Feature::Sqlite)
+            .expect("test config should allow sqlite");
+
+        let state_db = codex_state::StateRuntime::init(
+            home.path().to_path_buf(),
+            config.model_provider_id.clone(),
+            None,
+        )
+        .await
+        .expect("state db should initialize");
+        let thread_id = ThreadId::new();
+        let rollout_path = home.path().join("rollout.jsonl");
+        let builder = ThreadMetadataBuilder::new(
+            thread_id,
+            rollout_path.clone(),
+            Utc::now(),
+            SessionSource::Cli,
+        );
+        let items = vec![RolloutItem::EventMsg(EventMsg::AgentMessage(
+            AgentMessageEvent {
+                message: "assistant text".to_string(),
+                phase: None,
+            },
+        ))];
+
+        sync_thread_state_after_write(
+            Some(state_db.as_ref()),
+            rollout_path.as_path(),
+            Some(&builder),
+            items.as_slice(),
+            config.model_provider_id.as_str(),
+            None,
+        )
+        .await;
+
+        let thread = state_db
+            .get_thread(thread_id)
+            .await
+            .expect("thread should load after fallback")
+            .expect("thread should be inserted after fallback");
+        assert_eq!(thread.id, thread_id);
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn list_threads_db_disabled_does_not_skip_paginated_items() -> std::io::Result<()> {
         let home = TempDir::new().expect("temp dir");

codex-rs/core/src/state_db.rs

@@ -361,6 +361,7 @@ pub async fn reconcile_rollout(
             items,
             "reconcile_rollout",
             new_thread_memory_mode,
+            None,
         )
         .await;
         return;
@@ -491,6 +492,7 @@ pub async fn read_repair_rollout_path(
 }
 
 /// Apply rollout items incrementally to SQLite.
+#[allow(clippy::too_many_arguments)]
 pub async fn apply_rollout_items(
     context: Option<&codex_state::StateRuntime>,
     rollout_path: &Path,
@@ -499,6 +501,7 @@ pub async fn apply_rollout_items(
     items: &[RolloutItem],
     stage: &str,
     new_thread_memory_mode: Option<&str>,
+    updated_at_override: Option<DateTime<Utc>>,
 ) {
     let Some(ctx) = context else {
         return;
@@ -520,7 +523,13 @@ pub async fn apply_rollout_items(
     builder.rollout_path = rollout_path.to_path_buf();
     builder.cwd = normalize_cwd_for_state_db(&builder.cwd);
     if let Err(err) = ctx
-        .apply_rollout_items(&builder, items, None, new_thread_memory_mode)
+        .apply_rollout_items(
+            &builder,
+            items,
+            None,
+            new_thread_memory_mode,
+            updated_at_override,
+        )
         .await
     {
         warn!(
@@ -530,6 +539,26 @@ pub async fn apply_rollout_items(
     }
 }
 
+pub async fn touch_thread_updated_at(
+    context: Option<&codex_state::StateRuntime>,
+    thread_id: Option<ThreadId>,
+    updated_at: DateTime<Utc>,
+    stage: &str,
+) -> bool {
+    let Some(ctx) = context else {
+        return false;
+    };
+    let Some(thread_id) = thread_id else {
+        return false;
+    };
+    ctx.touch_thread_updated_at(thread_id, updated_at)
+        .await
+        .unwrap_or_else(|err| {
+            warn!("state db touch_thread_updated_at failed during {stage} for {thread_id}: {err}");
+            false
+        })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

codex-rs/state/src/extract.rs

@@ -29,6 +29,17 @@ pub fn apply_rollout_item(
     }
 }
 
+/// Return whether this rollout item can mutate thread metadata stored in SQLite.
+pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool {
+    match item {
+        RolloutItem::SessionMeta(_) | RolloutItem::TurnContext(_) => true,
+        RolloutItem::EventMsg(EventMsg::TokenCount(_) | EventMsg::UserMessage(_)) => true,
+        RolloutItem::EventMsg(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) => {
+            false
+        }
+    }
+}
+
 fn apply_session_meta_from_item(metadata: &mut ThreadMetadata, meta_line: &SessionMetaLine) {
     if metadata.id != meta_line.meta.id {
         // Ignore session_meta lines that don't match the canonical thread ID,

codex-rs/state/src/lib.rs

@@ -23,6 +23,7 @@ pub use runtime::StateRuntime;
 ///
 /// Most consumers should prefer [`StateRuntime`].
 pub use extract::apply_rollout_item;
+pub use extract::rollout_item_affects_thread_metadata;
 pub use model::AgentJob;
 pub use model::AgentJobCreateParams;
 pub use model::AgentJobItem;

codex-rs/state/src/runtime/threads.rs

@@ -278,6 +278,19 @@ ON CONFLICT(id) DO NOTHING
         Ok(result.rows_affected() > 0)
     }
 
+    pub async fn touch_thread_updated_at(
+        &self,
+        thread_id: ThreadId,
+        updated_at: DateTime<Utc>,
+    ) -> anyhow::Result<bool> {
+        let result = sqlx::query("UPDATE threads SET updated_at = ? WHERE id = ?")
+            .bind(datetime_to_epoch_seconds(updated_at))
+            .bind(thread_id.to_string())
+            .execute(self.pool.as_ref())
+            .await?;
+        Ok(result.rows_affected() > 0)
+    }
+
     pub async fn update_thread_git_info(
         &self,
         thread_id: ThreadId,
@@ -436,6 +449,7 @@ ON CONFLICT(thread_id, position) DO NOTHING
         items: &[RolloutItem],
         otel: Option<&OtelManager>,
         new_thread_memory_mode: Option<&str>,
+        updated_at_override: Option<DateTime<Utc>>,
     ) -> anyhow::Result<()> {
         if items.is_empty() {
             return Ok(());
@@ -451,7 +465,11 @@ ON CONFLICT(thread_id, position) DO NOTHING
         if let Some(existing_metadata) = existing_metadata.as_ref() {
             metadata.prefer_existing_git_info(existing_metadata);
         }
-        if let Some(updated_at) = file_modified_time_utc(builder.rollout_path.as_path()).await {
+        let updated_at = match updated_at_override {
+            Some(updated_at) => Some(updated_at),
+            None => file_modified_time_utc(builder.rollout_path.as_path()).await,
+        };
+        if let Some(updated_at) = updated_at {
             metadata.updated_at = updated_at;
         }
         // Keep the thread upsert before dynamic tools to satisfy the foreign key constraint:
@@ -649,6 +667,7 @@ mod tests {
     use super::*;
     use crate::runtime::test_support::test_thread_metadata;
     use crate::runtime::test_support::unique_temp_dir;
+    use codex_protocol::protocol::EventMsg;
     use codex_protocol::protocol::GitInfo;
     use codex_protocol::protocol::SessionMeta;
     use codex_protocol::protocol::SessionMetaLine;
@@ -735,7 +754,7 @@ mod tests {
         })];
 
         runtime
-            .apply_rollout_items(&builder, &items, None, None)
+            .apply_rollout_items(&builder, &items, None, None, None)
             .await
             .expect("apply_rollout_items should succeed");
 
@@ -793,7 +812,7 @@ mod tests {
         })];
 
         runtime
-            .apply_rollout_items(&builder, &items, None, None)
+            .apply_rollout_items(&builder, &items, None, None, None)
             .await
             .expect("apply_rollout_items should succeed");
 
@@ -947,4 +966,95 @@ mod tests {
         assert_eq!(persisted.git_branch, None);
         assert_eq!(persisted.git_origin_url, None);
     }
+
+    #[tokio::test]
+    async fn touch_thread_updated_at_updates_only_updated_at() {
+        let codex_home = unique_temp_dir();
+        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
+            .await
+            .expect("state db should initialize");
+        let thread_id =
+            ThreadId::from_string("00000000-0000-0000-0000-000000000791").expect("valid thread id");
+        let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone());
+        metadata.title = "original title".to_string();
+        metadata.first_user_message = Some("first-user-message".to_string());
+
+        runtime
+            .upsert_thread(&metadata)
+            .await
+            .expect("initial upsert should succeed");
+
+        let touched_at = DateTime::<Utc>::from_timestamp(1_700_001_111, 0).expect("timestamp");
+        let touched = runtime
+            .touch_thread_updated_at(thread_id, touched_at)
+            .await
+            .expect("touch should succeed");
+        assert!(touched);
+
+        let persisted = runtime
+            .get_thread(thread_id)
+            .await
+            .expect("thread should load")
+            .expect("thread should exist");
+        assert_eq!(persisted.updated_at, touched_at);
+        assert_eq!(persisted.title, "original title");
+        assert_eq!(
+            persisted.first_user_message.as_deref(),
+            Some("first-user-message")
+        );
+    }
+
+    #[tokio::test]
+    async fn apply_rollout_items_uses_override_updated_at_when_provided() {
+        let codex_home = unique_temp_dir();
+        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
+            .await
+            .expect("state db should initialize");
+        let thread_id =
+            ThreadId::from_string("00000000-0000-0000-0000-000000000792").expect("valid thread id");
+        let metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone());
+
+        runtime
+            .upsert_thread(&metadata)
+            .await
+            .expect("initial upsert should succeed");
+
+        let builder = ThreadMetadataBuilder::new(
+            thread_id,
+            metadata.rollout_path.clone(),
+            metadata.created_at,
+            SessionSource::Cli,
+        );
+        let items = vec![RolloutItem::EventMsg(EventMsg::TokenCount(
+            codex_protocol::protocol::TokenCountEvent {
+                info: Some(codex_protocol::protocol::TokenUsageInfo {
+                    total_token_usage: codex_protocol::protocol::TokenUsage {
+                        input_tokens: 0,
+                        cached_input_tokens: 0,
+                        output_tokens: 0,
+                        reasoning_output_tokens: 0,
+                        total_tokens: 321,
+                    },
+                    last_token_usage: codex_protocol::protocol::TokenUsage::default(),
+                    model_context_window: None,
+                }),
+                rate_limits: None,
+            },
+        ))];
+        let override_updated_at =
+            DateTime::<Utc>::from_timestamp(1_700_001_234, 0).expect("timestamp");
+
+        runtime
+            .apply_rollout_items(&builder, &items, None, None, Some(override_updated_at))
+            .await
+            .expect("apply_rollout_items should succeed");
+
+        let persisted = runtime
+            .get_thread(thread_id)
+            .await
+            .expect("thread should load")
+            .expect("thread should exist");
+        assert_eq!(persisted.tokens_used, 321);
+        assert_eq!(persisted.updated_at, override_updated_at);
+    }
 }