From c60cda34a3ba6975f698d63da9e8ca8a1616f0c9 Mon Sep 17 00:00:00 2001 From: Dmitriy Kovalenko Date: Mon, 2 Nov 2026 23:53:56 +0000 Subject: [PATCH] Make observed thread metadata writes atomic and retryable (#59016) ## Why A failed memory-mode update could leave other metadata changes applied. Failed observed metadata writes for paginated threads also need to remain pending so persistence can retry them. ## What changed - Apply observed metadata, memory mode, and optional recency advancement in one SQLite upsert, preserving monotonic recency. - Require successful SQLite writes for paginated thread metadata when a state runtime is configured, retaining failed writes for retry. - Use SQLite busy-retry helpers for read-state writes and immediate transactions. ## Testing Add regression tests that reject SQLite writes to verify pending metadata survives a failed persist and a rejected memory-mode update leaves existing metadata unchanged. Both tests verify a subsequent retry succeeds. GitOrigin-RevId: 652bf6ebee0197976852aaf9f096b907697fd590 --- .../state/src/runtime/thread_read_state.rs | 6 +- codex-rs/state/src/runtime/threads.rs | 45 ++++++- .../local/pending_thread_metadata_tests.rs | 125 ++++++++++++++++++ .../src/local/update_thread_metadata.rs | 35 ++--- 4 files changed, 182 insertions(+), 29 deletions(-) diff --git a/codex-rs/state/src/runtime/thread_read_state.rs b/codex-rs/state/src/runtime/thread_read_state.rs index 31bb985a0a..bf48186488 100644 --- a/codex-rs/state/src/runtime/thread_read_state.rs +++ b/codex-rs/state/src/runtime/thread_read_state.rs @@ -12,6 +12,8 @@ use sqlx::SqliteConnection; use uuid::Uuid; use super::StateRuntime; +use crate::sqlite::SqlitePoolWriteExt; +use crate::sqlite::SqliteWriteQueryExt; #[derive(Debug, Clone, PartialEq, Eq)] pub struct ThreadReadState { @@ -63,7 +65,7 @@ impl StateRuntime { .bind(Uuid::new_v4().to_string()) .bind(turn_id) .bind(id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? .rows_affected() != 0; Ok(changed) @@ -77,7 +79,7 @@ impl StateRuntime { expected_revision: &str, operation: ReadStateOperation, ) -> anyhow::Result { - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; let Some(before) = read_states(&mut tx, &[id]).await?.remove(&id) else { return Ok(ReadStateUpdate::Unavailable); }; diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 8635dd4e8c..df59ba60a7 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -932,11 +932,41 @@ WHERE id = ? Ok(result.rows_affected() > 0) } + /// Apply observed metadata, memory preference, and recency in one commit. + pub async fn upsert_thread_with_observed_metadata( + &self, + metadata: &crate::ThreadMetadata, + memory_mode: Option<&str>, + recency_at: Option>, + ) -> anyhow::Result<()> { + self.upsert_thread_metadata(metadata, memory_mode, memory_mode, recency_at) + .await + } + async fn upsert_thread_with_creation_memory_mode( &self, metadata: &crate::ThreadMetadata, creation_memory_mode: Option<&str>, ) -> anyhow::Result<()> { + self.upsert_thread_metadata( + metadata, + creation_memory_mode, + /*memory_mode*/ None, + /*recency_at*/ None, + ) + .await + } + + async fn upsert_thread_metadata( + &self, + metadata: &crate::ThreadMetadata, + creation_memory_mode: Option<&str>, + memory_mode: Option<&str>, + recency_at: Option>, + ) -> anyhow::Result<()> { + let recency_at = recency_at + .map(|at| self.allocate_thread_recency_at(at)) + .transpose()?; let updated_at = self.allocate_thread_updated_at(metadata.updated_at)?; let insert_recency_at = self.allocate_thread_recency_at(metadata.recency_at)?; let preview = metadata_preview(metadata); @@ -994,10 +1024,12 @@ ON CONFLICT(id) DO UPDATE SET rollout_path = excluded.rollout_path, created_at = excluded.created_at, updated_at = excluded.updated_at, - recency_at = threads.recency_at, + recency_at = CASE WHEN ? IS NULL THEN threads.recency_at + ELSE MAX(?, MAX(?, threads.recency_at_ms + 1) / 1000) END, created_at_ms = excluded.created_at_ms, updated_at_ms = excluded.updated_at_ms, - recency_at_ms = threads.recency_at_ms, + recency_at_ms = CASE WHEN ? IS NULL THEN threads.recency_at_ms + ELSE MAX(?, threads.recency_at_ms + 1) END, source = excluded.source, originator = COALESCE(threads.originator, excluded.originator), creator_user_id = COALESCE(threads.creator_user_id, excluded.creator_user_id), @@ -1026,7 +1058,8 @@ ON CONFLICT(id) DO UPDATE SET archived_at = excluded.archived_at, git_sha = COALESCE(threads.git_sha, excluded.git_sha), git_branch = COALESCE(threads.git_branch, excluded.git_branch), - git_origin_url = COALESCE(threads.git_origin_url, excluded.git_origin_url) + git_origin_url = COALESCE(threads.git_origin_url, excluded.git_origin_url), + memory_mode = COALESCE(?, threads.memory_mode) "#, ) .bind(metadata.id.to_string()) @@ -1080,6 +1113,12 @@ ON CONFLICT(id) DO UPDATE SET .bind(metadata.project_id.as_deref()) .bind(metadata.daybreak_enabled) .bind(metadata.attached_thread_parent_id.map(|id| id.to_string())) + .bind(recency_at.map(datetime_to_epoch_seconds)) + .bind(recency_at.map(datetime_to_epoch_seconds)) + .bind(recency_at.map(datetime_to_epoch_millis)) + .bind(recency_at.map(datetime_to_epoch_millis)) + .bind(recency_at.map(datetime_to_epoch_millis)) + .bind(memory_mode) .execute_with_busy_retry(self.pool.as_ref()) .await?; self.insert_thread_spawn_edge_from_source_if_absent(metadata.id, metadata.source.as_str()) diff --git a/codex-rs/thread-store/src/local/pending_thread_metadata_tests.rs b/codex-rs/thread-store/src/local/pending_thread_metadata_tests.rs index 8c6856588c..a81eb389a7 100644 --- a/codex-rs/thread-store/src/local/pending_thread_metadata_tests.rs +++ b/codex-rs/thread-store/src/local/pending_thread_metadata_tests.rs @@ -16,10 +16,135 @@ use tempfile::TempDir; use super::*; use crate::CreateThreadParams; use crate::LiveThread; +use crate::PersistContext; use crate::ThreadPersistenceMetadata; use crate::UpdateThreadMetadataParams; use crate::local::test_support::test_config; +#[tokio::test] +async fn paginated_metadata_write_failure_stays_pending_until_retry() { + let (_home, store, runtime) = store_with_runtime().await; + let thread_id = ThreadId::new(); + let mut params = create_thread_params(thread_id); + params.history_mode = ThreadHistoryMode::Paginated; + let live_thread = LiveThread::create(store.clone(), params) + .await + .expect("create live thread"); + let pool = store + .config + .sqlite + .open_read_write_pool(&store.config.sqlite.state_db_path()) + .await + .expect("open SQL observer"); + sqlx::query( + "CREATE TRIGGER reject_metadata BEFORE INSERT ON threads BEGIN SELECT RAISE(ABORT, 'metadata write rejected'); END", + ) + .execute(&pool) + .await + .expect("reject index writes"); + + let error = live_thread + .persist(PersistContext::Standard) + .await + .expect_err("failed index writes must remain pending"); + assert!(error.to_string().contains("metadata write rejected")); + assert_eq!( + runtime.get_thread(thread_id).await.expect("read index"), + None + ); + + sqlx::query("DROP TRIGGER reject_metadata") + .execute(&pool) + .await + .expect("allow index writes"); + live_thread + .persist(PersistContext::Standard) + .await + .expect("retry pending metadata"); + let metadata = runtime + .get_thread(thread_id) + .await + .expect("read index") + .expect("pending metadata was indexed"); + assert_eq!( + (metadata.id, metadata.history_mode, metadata.model_provider), + ( + thread_id, + ThreadHistoryMode::Paginated, + "test-provider".to_string() + ) + ); + live_thread + .shutdown() + .await + .expect("shutdown indexed thread"); +} + +#[tokio::test] +async fn metadata_patch_failure_does_not_partially_apply() { + let (_home, store, runtime) = store_with_runtime().await; + let thread_id = ThreadId::new(); + let mut params = create_thread_params(thread_id); + params.history_mode = ThreadHistoryMode::Paginated; + let live_thread = LiveThread::create(store.clone(), params) + .await + .expect("create thread"); + live_thread + .persist(PersistContext::Standard) + .await + .expect("persist thread"); + let before = runtime.get_thread(thread_id).await.expect("read metadata"); + let pool = store + .config + .sqlite + .open_read_write_pool(&store.config.sqlite.state_db_path()) + .await + .expect("open SQL observer"); + sqlx::query("CREATE TRIGGER reject_memory_mode BEFORE UPDATE OF memory_mode ON threads BEGIN SELECT RAISE(ABORT, 'memory mode rejected'); END") + .execute(&pool).await.expect("reject memory preference update"); + let patch = ThreadMetadataPatch { + model: Some("changed-model".to_string()), + memory_mode: Some(ThreadMemoryMode::Disabled), + ..Default::default() + }; + live_thread + .update_metadata(patch.clone(), /*include_archived*/ false) + .await + .expect_err("reject the entire patch"); + assert_eq!( + runtime.get_thread(thread_id).await.expect("read metadata"), + before + ); + + sqlx::query("DROP TRIGGER reject_memory_mode") + .execute(&pool) + .await + .expect("allow update"); + live_thread + .update_metadata(patch, /*include_archived*/ false) + .await + .expect("retry patch"); + let metadata = runtime + .get_thread(thread_id) + .await + .expect("read metadata") + .expect("thread"); + assert_eq!( + ( + metadata.model, + runtime + .get_thread_memory_mode(thread_id) + .await + .expect("memory mode") + ), + ( + Some("changed-model".to_string()), + Some("disabled".to_string()) + ) + ); + live_thread.shutdown().await.expect("shutdown thread"); +} + #[tokio::test] async fn naming_empty_paginated_thread_materializes_rollout() { let (_home, store, runtime) = store_with_runtime().await; diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index 33fd5473ae..2aab5d9614 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -73,10 +73,12 @@ pub(super) async fn update_thread_metadata( let requires_rollout_compat = staged_requires_rollout_compat || requires_rollout_compatibility_update(&patch); let has_explicit_metadata = patch.name.is_some() || requires_rollout_compat; - let history_mode = if has_explicit_metadata { + // Retain failed observed writes when indexing is available, while preserving callers + // that intentionally run without a state runtime. + let history_mode = if has_explicit_metadata || store.state_db.is_some() { match live_writer::live_writer_parts(store, thread_id).await { Ok((_recorder, _rollout_id, history_mode)) => Some(history_mode), - Err(ThreadStoreError::ThreadNotFound { .. }) => Some( + Err(ThreadStoreError::ThreadNotFound { .. }) if has_explicit_metadata => Some( read_thread::read_thread( store, ReadThreadParams { @@ -88,6 +90,7 @@ pub(super) async fn update_thread_metadata( .await? .history_mode, ), + Err(ThreadStoreError::ThreadNotFound { .. }) => None, Err(err) => return Err(err), } } else { @@ -522,13 +525,17 @@ async fn apply_metadata_update( if let Some(daybreak_enabled) = patch.daybreak_enabled { metadata.daybreak_enabled = Some(daybreak_enabled); } - let upsert_result = state_db.upsert_thread(&metadata).await; + let memory_mode = patch.memory_mode.map(memory_mode_as_str); + let recency_at = existing.as_ref().and(advance_recency_at); + let upsert_result = state_db + .upsert_thread_with_observed_metadata(&metadata, memory_mode, recency_at) + .await; if existing.is_none() && metadata.project_id.is_some() && matches!(&upsert_result, Err(err) if err.to_string().contains("FOREIGN KEY constraint failed")) { metadata.project_id = None; - state_db.upsert_thread(&metadata).await.map_err(|err| { + state_db.upsert_thread_with_observed_metadata(&metadata, memory_mode, recency_at).await.map_err(|err| { ThreadStoreError::Internal { message: format!("failed to update thread metadata for {thread_id}: {err}"), } @@ -591,26 +598,6 @@ async fn apply_metadata_update( }); } } - if existing.is_some() - && let Some(recency_at) = advance_recency_at - { - state_db - .touch_thread_recency_at(thread_id, recency_at) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!( - "failed to advance thread recency_at for {thread_id}: {err}" - ), - })?; - } - if let Some(memory_mode) = patch.memory_mode { - state_db - .set_thread_memory_mode(thread_id, memory_mode_as_str(memory_mode)) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to update memory mode for {thread_id}: {err}"), - })?; - } if let Some(daybreak_enabled) = patch.daybreak_enabled && !state_db .set_thread_daybreak_enabled(thread_id, daybreak_enabled)