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
This commit is contained in:
Dmitriy Kovalenko
2026-11-03 00:16:13 +00:00
committed by copyberry
parent 9e77fb4ad5
commit c60cda34a3
4 changed files with 182 additions and 29 deletions
@@ -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;
@@ -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)