From e0012eb0bc210a01b86b79bf79675089bfc9edad Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Sat, 25 Jul 2026 19:28:03 -0400 Subject: [PATCH 1/3] fix: harden reviewed trust boundaries --- changelog.d/review-hardening.fixed.md | 1 + n00n-agent/src/agent/compaction.rs | 6 +++- n00n-agent/src/mcp/oauth/discovery.rs | 42 +++++++++++++++++-------- n00n-lua/src/plugin_permissions.rs | 13 ++++---- n00n-ui/src/storage_writer.rs | 45 +++++++++++++++++++++++++-- 5 files changed, 84 insertions(+), 23 deletions(-) create mode 100644 changelog.d/review-hardening.fixed.md diff --git a/changelog.d/review-hardening.fixed.md b/changelog.d/review-hardening.fixed.md new file mode 100644 index 000000000..626cd8b18 --- /dev/null +++ b/changelog.d/review-hardening.fixed.md @@ -0,0 +1 @@ +Hardened project plugin permissions and OAuth endpoint validation, bounded compaction overflow retries, and retained failed session snapshots for retry. diff --git a/n00n-agent/src/agent/compaction.rs b/n00n-agent/src/agent/compaction.rs index 17034cb3c..27c986a22 100644 --- a/n00n-agent/src/agent/compaction.rs +++ b/n00n-agent/src/agent/compaction.rs @@ -122,8 +122,12 @@ pub(super) async fn compact_history( .await { Ok(response) => break response, - Err(e) if e.is_context_overflow() => { + Err(error) if error.is_context_overflow() => { + let previous_len = compaction_history.len(); truncate_oldest_round(&mut compaction_history); + if compaction_history.len() == previous_len { + return Err(error); + } } Err(e) => return Err(e), } diff --git a/n00n-agent/src/mcp/oauth/discovery.rs b/n00n-agent/src/mcp/oauth/discovery.rs index 684d125ff..a5235e69c 100644 --- a/n00n-agent/src/mcp/oauth/discovery.rs +++ b/n00n-agent/src/mcp/oauth/discovery.rs @@ -1,4 +1,5 @@ use std::io::Read; +use std::net::IpAddr; use isahc::HttpClient; use isahc::http::Request; @@ -96,16 +97,26 @@ pub(super) fn server_origin(server_url: &str) -> String { } fn validate_endpoint_url(url: &str) -> Result<(), OAuthError> { - if url.starts_with("https://") - || url.starts_with("http://127.0.0.1") - || url.starts_with("http://localhost") - { - Ok(()) - } else { - Err(OAuthError::Other(format!( - "endpoint URL must use HTTPS: {url}" - ))) + let uri = url + .parse::() + .map_err(|error| OAuthError::Other(format!("invalid endpoint URL '{url}': {error}")))?; + let scheme = uri.scheme_str(); + let host = uri.host(); + let secure = scheme == Some("https") && host.is_some(); + let loopback = scheme == Some("http") + && host.is_some_and(|host| { + let host = host.trim_start_matches('[').trim_end_matches(']'); + host.eq_ignore_ascii_case("localhost") + || host + .parse::() + .is_ok_and(|address| address.is_loopback()) + }); + if secure || loopback { + return Ok(()); } + Err(OAuthError::Other(format!( + "endpoint URL must use HTTPS or an exact loopback host: {url}" + ))) } /// Validates that all required endpoints in auth server metadata use HTTPS. @@ -278,10 +289,15 @@ mod tests { assert_eq!(well_known_url(base, name), expected); } - #[test_case("https://example.com/token", true ; "https_valid")] - #[test_case("http://127.0.0.1:8080/token", true ; "localhost_valid")] - #[test_case("http://example.com/token", false ; "http_rejected")] - #[test_case("ftp://example.com/token", false ; "ftp_rejected")] + #[test_case("https://example.com/token", true ; "https_valid")] + #[test_case("http://127.0.0.1:8080/token", true ; "ipv4_loopback_valid")] + #[test_case("http://[::1]:8080/token", true ; "ipv6_loopback_valid")] + #[test_case("http://localhost:8080/token", true ; "localhost_valid")] + #[test_case("http://localhost.attacker.invalid/token", false ; "localhost_prefix_rejected")] + #[test_case("http://127.0.0.1.attacker.invalid/token", false ; "ipv4_prefix_rejected")] + #[test_case("http://example.com/token", false ; "http_rejected")] + #[test_case("ftp://example.com/token", false ; "ftp_rejected")] + #[test_case("not a url", false ; "malformed_rejected")] fn endpoint_url_validation(url: &str, should_pass: bool) { assert_eq!(validate_endpoint_url(url).is_ok(), should_pass); } diff --git a/n00n-lua/src/plugin_permissions.rs b/n00n-lua/src/plugin_permissions.rs index 55db21a6a..c6de9b3a8 100644 --- a/n00n-lua/src/plugin_permissions.rs +++ b/n00n-lua/src/plugin_permissions.rs @@ -65,14 +65,13 @@ impl PluginPermissions { } pub fn from_manifest(manifest: &toml::Value) -> Self { - const DEFAULT_PERMISSION: bool = true; let perms = manifest.get("permissions"); - let mut allowed = [true; 5]; + let mut allowed = [false; 5]; for perm in Permission::ALL { allowed[perm as usize] = perms .and_then(|p| p.get(perm.manifest_key())) .and_then(toml::Value::as_bool) - .unwrap_or_else(|| DEFAULT_PERMISSION); + == Some(true); } Self { allowed } } @@ -189,10 +188,10 @@ mod tests { .unwrap(); let p = PluginPermissions::from_manifest(&val); assert!(!p.is_allowed(Permission::FsRead)); - assert!(p.is_allowed(Permission::FsWrite)); + assert!(!p.is_allowed(Permission::FsWrite)); assert!(!p.is_allowed(Permission::Net)); - assert!(p.is_allowed(Permission::Run)); - assert!(p.is_allowed(Permission::Env)); + assert!(!p.is_allowed(Permission::Run)); + assert!(!p.is_allowed(Permission::Env)); } #[test] @@ -200,7 +199,7 @@ mod tests { let val: toml::Value = toml::from_str("[package]\nname = \"test\"").unwrap(); let p = PluginPermissions::from_manifest(&val); for perm in Permission::ALL { - assert!(p.is_allowed(perm), "{perm} should default to allowed"); + assert!(!p.is_allowed(perm), "{perm} should default to denied"); } } diff --git a/n00n-ui/src/storage_writer.rs b/n00n-ui/src/storage_writer.rs index 570644c44..817878ad4 100644 --- a/n00n-ui/src/storage_writer.rs +++ b/n00n-ui/src/storage_writer.rs @@ -235,13 +235,21 @@ fn flush( return true; } }; - for snapshot in batch.into_values() { + let mut failed = false; + for (id, snapshot) in batch { if let Err(error) = write_session(&sessions_dir, logs, durable_revisions, &snapshot.session) { warn!(error = %error, id = %snapshot.session.id, "session write failed"); + let replace = pending_guard + .get(&id) + .is_none_or(|current| current.revision < snapshot.revision); + if replace { + pending_guard.insert(id, snapshot); + } + failed = true; } } - false + failed } fn persist_session( @@ -442,6 +450,39 @@ mod tests { writer.shutdown(DRAIN_TIMEOUT); } + #[test] + fn failed_snapshot_remains_pending_for_retry() { + let (_tmp, dir) = state_dir(); + let session = AppSession::new("test-model", "/tmp/retry"); + let id = session.id; + let revision = session.meta.revision; + let sessions_dir = dir.ensure_subdir(SESSIONS_DIR).unwrap(); + let blocked_log = sessions_dir.join(format!("{id}.jsonl")); + std::fs::create_dir(&blocked_log).unwrap(); + + let pending: Pending = Arc::default(); + lock(&pending).insert( + id, + PendingSnapshot { + revision, + session: Box::new(session), + }, + ); + let mut logs = HashMap::new(); + let mut durable_revisions = HashMap::new(); + + assert!(flush(&pending, &mut logs, &mut durable_revisions, &dir)); + assert_eq!( + lock(&pending).get(&id).map(|item| item.revision), + Some(revision) + ); + + std::fs::remove_dir(&blocked_log).unwrap(); + assert!(!flush(&pending, &mut logs, &mut durable_revisions, &dir)); + assert!(lock(&pending).is_empty()); + assert!(AppSession::load(id, &dir).is_ok()); + } + #[test] fn delete_discards_pending_snapshot() { let (_tmp, dir) = state_dir(); From c37ac0ae1a58534fbe732009bfcf8e59a72e0913 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Sat, 25 Jul 2026 20:10:23 -0400 Subject: [PATCH 2/3] fix: isolate session persistence retries --- n00n-agent/src/agent/compaction.rs | 38 ++++- n00n-lua/src/plugin_permissions.rs | 4 +- n00n-ui/src/storage_writer.rs | 244 ++++++++++++++++++++++------- 3 files changed, 222 insertions(+), 64 deletions(-) diff --git a/n00n-agent/src/agent/compaction.rs b/n00n-agent/src/agent/compaction.rs index 27c986a22..1d3f6b92a 100644 --- a/n00n-agent/src/agent/compaction.rs +++ b/n00n-agent/src/agent/compaction.rs @@ -328,11 +328,15 @@ mod tests { use crate::AgentConfig; struct MockProvider { - responses: Mutex>, + responses: Mutex>>, } impl MockProvider { fn new(responses: Vec) -> Self { + Self::with_results(responses.into_iter().map(Ok).collect()) + } + + fn with_results(responses: Vec>) -> Self { Self { responses: Mutex::new(responses), } @@ -353,7 +357,7 @@ mod tests { Box::pin(async { let mut responses = self.responses.lock().unwrap(); assert!(!responses.is_empty(), "MockProvider: no more responses"); - Ok(responses.remove(0)) + responses.remove(0) }) } @@ -458,6 +462,36 @@ mod tests { }); } + #[test] + fn compact_history_returns_overflow_when_nothing_can_be_truncated() { + smol::block_on(async { + let provider = MockProvider::with_results(vec![Err(AgentError::Api { + status: 413, + message: "context overflow".into(), + })]); + let model = default_model(); + let (raw_tx, _rx) = flume::unbounded(); + let mut history = History::new(Vec::new()); + + let error = compact_history( + &provider, + &model, + &mut history, + &EventSender::new(raw_tx, 0), + &CancelToken::none(), + CompactionTrigger::Manual, + None, + std::path::Path::new("/tmp"), + None, + ) + .await + .unwrap_err(); + + assert!(error.is_context_overflow()); + assert!(history.as_slice().is_empty()); + }); + } + #[test_case(159_999, 0, 0, 0, 200_000, false ; "below_threshold")] #[test_case(160_000, 0, 0, 0, 200_000, true ; "at_threshold")] #[test_case(100, 0, 0, 0, 100, true ; "tiny_context_window")] diff --git a/n00n-lua/src/plugin_permissions.rs b/n00n-lua/src/plugin_permissions.rs index c6de9b3a8..1835b8906 100644 --- a/n00n-lua/src/plugin_permissions.rs +++ b/n00n-lua/src/plugin_permissions.rs @@ -181,13 +181,13 @@ mod tests { let val: toml::Value = toml::from_str( r" [permissions] - fs_read = false + fs_read = true net = false ", ) .unwrap(); let p = PluginPermissions::from_manifest(&val); - assert!(!p.is_allowed(Permission::FsRead)); + assert!(p.is_allowed(Permission::FsRead)); assert!(!p.is_allowed(Permission::FsWrite)); assert!(!p.is_allowed(Permission::Net)); assert!(!p.is_allowed(Permission::Run)); diff --git a/n00n-ui/src/storage_writer.rs b/n00n-ui/src/storage_writer.rs index 817878ad4..402f55e43 100644 --- a/n00n-ui/src/storage_writer.rs +++ b/n00n-ui/src/storage_writer.rs @@ -24,7 +24,14 @@ struct PendingSnapshot { session: Box, } +#[derive(Clone, Copy)] +struct RetryState { + revision: u64, + attempts: u32, +} + type Pending = Arc>>; +type FailedSnapshots = HashMap; type DeleteCallback = Box) + Send>; type PersistCallback = Box) + Send>; @@ -62,70 +69,40 @@ impl StorageWriter { .spawn(move || { let mut logs: HashMap = HashMap::new(); let mut durable_revisions: HashMap = HashMap::new(); - let mut retry_count: u32 = 0; + let mut retry_states: HashMap = HashMap::new(); loop { - let op = if retry_count > 0 { + let op = if lock(&writer_pending).is_empty() { + ops_rx.recv().ok() + } else { match ops_rx.recv_timeout(RETRY_DELAY) { Ok(op) => Some(op), Err(flume::RecvTimeoutError::Timeout) => Some(Op::Flush), Err(flume::RecvTimeoutError::Disconnected) => None, } - } else { - ops_rx.recv().ok() }; let Some(op) = op else { break }; match op { Op::Flush => { - let failed = flush(&writer_pending, &mut logs, &mut durable_revisions, &dir); - if failed { - retry_count += 1; - if retry_count >= MAX_RETRY_ATTEMPTS { - warn!( - retry_count, - "storage writer exhausted retry attempts, dropping pending snapshots" - ); - lock(&writer_pending).clear(); - retry_count = 0; - } - } else { - retry_count = 0; - } + let failed = + flush(&writer_pending, &mut logs, &mut durable_revisions, &dir); + update_retry_states(&writer_pending, &mut retry_states, &failed); } Op::Persist { session, done } => { - let failed = flush(&writer_pending, &mut logs, &mut durable_revisions, &dir); - if failed { - retry_count += 1; - if retry_count >= MAX_RETRY_ATTEMPTS { - warn!( - retry_count, - id = %session.id, - "storage writer exhausted retry attempts, dropping snapshot" - ); - lock(&writer_pending).remove(&session.id); - retry_count = 0; - done(Err(retry_exhausted())); - } else { - done(persist_session( - &writer_pending, - &mut logs, - &mut durable_revisions, - &dir, - &session, - )); - } - } else { - retry_count = 0; - done(persist_session( - &writer_pending, - &mut logs, - &mut durable_revisions, - &dir, - &session, - )); - } + let failed = + flush(&writer_pending, &mut logs, &mut durable_revisions, &dir); + let result = persist_session( + &writer_pending, + &mut logs, + &mut durable_revisions, + &dir, + &session, + ); + update_retry_states(&writer_pending, &mut retry_states, &failed); + done(result); } Op::Delete { id, done } => { lock(&writer_pending).remove(&id); + retry_states.remove(&id); logs.remove(&id); done(AppSession::delete(id, &dir)); } @@ -204,9 +181,51 @@ fn writer_gone() -> SessionError { n00n_storage::StorageError::Io(io::Error::other("storage writer unavailable")).into() } -fn retry_exhausted() -> SessionError { - n00n_storage::StorageError::Io(io::Error::other("storage writer exhausted retry attempts")) - .into() +fn update_retry_states( + pending: &Pending, + retry_states: &mut HashMap, + failed: &FailedSnapshots, +) { + let mut pending_guard = lock(pending); + retry_states.retain(|id, _| { + failed.get(id).is_some_and(|revision| { + pending_guard + .get(id) + .is_some_and(|snapshot| snapshot.revision == *revision) + }) + }); + let mut exhausted = Vec::new(); + for (&id, &revision) in failed { + if pending_guard + .get(&id) + .is_none_or(|snapshot| snapshot.revision != revision) + { + continue; + } + let state = retry_states.entry(id).or_insert(RetryState { + revision, + attempts: 0, + }); + if state.revision != revision { + state.revision = revision; + state.attempts = 0; + } + state.attempts += 1; + if state.attempts >= MAX_RETRY_ATTEMPTS { + exhausted.push((id, revision)); + } + } + + for (id, revision) in exhausted { + pending_guard.remove(&id); + warn!( + retry_count = MAX_RETRY_ATTEMPTS, + %id, + revision, + "storage writer exhausted retry attempts, dropping pending snapshot" + ); + retry_states.remove(&id); + } } fn flush( @@ -214,16 +233,20 @@ fn flush( logs: &mut HashMap, durable_revisions: &mut HashMap, dir: &StateDir, -) -> bool { +) -> FailedSnapshots { let mut pending_guard = lock(pending); let batch = mem::take(&mut *pending_guard); if batch.is_empty() { - return false; + return HashMap::new(); } let sessions_dir = match dir.ensure_subdir(SESSIONS_DIR) { Ok(d) => d, Err(e) => { warn!(error = %e, "failed to ensure sessions dir"); + let failed = batch + .iter() + .map(|(&id, snapshot)| (id, snapshot.revision)) + .collect(); for (id, snapshot) in batch { let replace = pending_guard .get(&id) @@ -232,10 +255,10 @@ fn flush( pending_guard.insert(id, snapshot); } } - return true; + return failed; } }; - let mut failed = false; + let mut failed = HashMap::new(); for (id, snapshot) in batch { if let Err(error) = write_session(&sessions_dir, logs, durable_revisions, &snapshot.session) { @@ -243,10 +266,10 @@ fn flush( let replace = pending_guard .get(&id) .is_none_or(|current| current.revision < snapshot.revision); + failed.insert(id, snapshot.revision); if replace { pending_guard.insert(id, snapshot); } - failed = true; } } failed @@ -264,7 +287,7 @@ fn persist_session( .get(&session.id) .is_some_and(|snapshot| snapshot.revision > session.meta.revision) { - let snapshot = pending_guard.remove(&session.id).ok_or_else(|| { + let snapshot = pending_guard.get(&session.id).ok_or_else(|| { n00n_storage::StorageError::Io(io::Error::other("pending snapshot disappeared")) })?; write_session( @@ -273,6 +296,7 @@ fn persist_session( durable_revisions, &snapshot.session, )?; + pending_guard.remove(&session.id); } drop(pending_guard); write_session_if_newer(logs, durable_revisions, dir, session) @@ -471,18 +495,118 @@ mod tests { let mut logs = HashMap::new(); let mut durable_revisions = HashMap::new(); - assert!(flush(&pending, &mut logs, &mut durable_revisions, &dir)); + assert!(!flush(&pending, &mut logs, &mut durable_revisions, &dir).is_empty()); assert_eq!( lock(&pending).get(&id).map(|item| item.revision), Some(revision) ); std::fs::remove_dir(&blocked_log).unwrap(); - assert!(!flush(&pending, &mut logs, &mut durable_revisions, &dir)); + assert!(flush(&pending, &mut logs, &mut durable_revisions, &dir).is_empty()); assert!(lock(&pending).is_empty()); assert!(AppSession::load(id, &dir).is_ok()); } + #[test] + fn persist_failure_retains_newer_pending_snapshot() { + let (_tmp, dir) = state_dir(); + let mut older = AppSession::new("test-model", "/tmp/persist-retry"); + older.meta.revision = 1; + let id = older.id; + let mut newer = older.clone(); + newer.meta.revision = 2; + let sessions_dir = dir.ensure_subdir(SESSIONS_DIR).unwrap(); + std::fs::create_dir(sessions_dir.join(format!("{id}.jsonl"))).unwrap(); + let pending: Pending = Arc::default(); + lock(&pending).insert( + id, + PendingSnapshot { + revision: newer.meta.revision, + session: Box::new(newer), + }, + ); + + let result = persist_session( + &pending, + &mut HashMap::new(), + &mut HashMap::new(), + &dir, + &older, + ); + + assert!(result.is_err()); + assert_eq!( + lock(&pending).get(&id).map(|snapshot| snapshot.revision), + Some(2) + ); + } + + #[test] + fn exhausted_retry_does_not_drop_newer_snapshot() { + let pending: Pending = Arc::default(); + let mut older = AppSession::new("test-model", "/tmp/older"); + older.meta.revision = 1; + let id = older.id; + let mut newer = older.clone(); + newer.meta.revision = 2; + lock(&pending).insert( + id, + PendingSnapshot { + revision: newer.meta.revision, + session: Box::new(newer), + }, + ); + + let failed = HashMap::from([(id, older.meta.revision)]); + let mut retry_states = HashMap::new(); + for _ in 0..MAX_RETRY_ATTEMPTS { + update_retry_states(&pending, &mut retry_states, &failed); + } + + assert_eq!( + lock(&pending).get(&id).map(|snapshot| snapshot.revision), + Some(2) + ); + assert!(!retry_states.contains_key(&id)); + } + + #[test] + fn retry_attempts_reset_for_newer_revision() { + let pending: Pending = Arc::default(); + let mut session = AppSession::new("test-model", "/tmp/revisions"); + session.meta.revision = 1; + let id = session.id; + let mut newer = session.clone(); + newer.meta.revision = 2; + lock(&pending).insert( + id, + PendingSnapshot { + revision: 1, + session: Box::new(session), + }, + ); + let mut retry_states = HashMap::new(); + let old_failure = HashMap::from([(id, 1)]); + for _ in 1..MAX_RETRY_ATTEMPTS { + update_retry_states(&pending, &mut retry_states, &old_failure); + } + lock(&pending).insert( + id, + PendingSnapshot { + revision: 2, + session: Box::new(newer), + }, + ); + + let new_failure = HashMap::from([(id, 2)]); + update_retry_states(&pending, &mut retry_states, &new_failure); + + let state = &retry_states[&id]; + assert_eq!(state.revision, 2); + assert_eq!(state.attempts, 1); + assert!(lock(&pending).contains_key(&id)); + } + #[test] fn delete_discards_pending_snapshot() { let (_tmp, dir) = state_dir(); From e9078a6fdf1f46b782244eb79254f92bee94badd Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Sat, 25 Jul 2026 20:23:01 -0400 Subject: [PATCH 3/3] fix: reset retries for replaced snapshots --- n00n-agent/src/mcp/oauth/discovery.rs | 4 +- n00n-lua/src/plugin_permissions.rs | 3 +- n00n-ui/src/storage_writer.rs | 192 +++++++++----------------- 3 files changed, 67 insertions(+), 132 deletions(-) diff --git a/n00n-agent/src/mcp/oauth/discovery.rs b/n00n-agent/src/mcp/oauth/discovery.rs index a5235e69c..c19416f06 100644 --- a/n00n-agent/src/mcp/oauth/discovery.rs +++ b/n00n-agent/src/mcp/oauth/discovery.rs @@ -2,7 +2,7 @@ use std::io::Read; use std::net::IpAddr; use isahc::HttpClient; -use isahc::http::Request; +use isahc::http::{Request, Uri}; use serde::Deserialize; use super::OAuthError; @@ -98,7 +98,7 @@ pub(super) fn server_origin(server_url: &str) -> String { fn validate_endpoint_url(url: &str) -> Result<(), OAuthError> { let uri = url - .parse::() + .parse::() .map_err(|error| OAuthError::Other(format!("invalid endpoint URL '{url}': {error}")))?; let scheme = uri.scheme_str(); let host = uri.host(); diff --git a/n00n-lua/src/plugin_permissions.rs b/n00n-lua/src/plugin_permissions.rs index 1835b8906..b86fa2b72 100644 --- a/n00n-lua/src/plugin_permissions.rs +++ b/n00n-lua/src/plugin_permissions.rs @@ -2,6 +2,7 @@ use std::fmt; use std::path::Path; use mlua::{Error as LuaError, Function, IntoLuaMulti, Lua, Result as LuaResult}; +use toml::Value; use tracing::warn; const MANIFEST_FILE: &str = "plugin.toml"; @@ -70,7 +71,7 @@ impl PluginPermissions { for perm in Permission::ALL { allowed[perm as usize] = perms .and_then(|p| p.get(perm.manifest_key())) - .and_then(toml::Value::as_bool) + .and_then(Value::as_bool) == Some(true); } Self { allowed } diff --git a/n00n-ui/src/storage_writer.rs b/n00n-ui/src/storage_writer.rs index 402f55e43..d4472d3e6 100644 --- a/n00n-ui/src/storage_writer.rs +++ b/n00n-ui/src/storage_writer.rs @@ -21,17 +21,11 @@ use crate::AppSession; struct PendingSnapshot { revision: u64, + retry_attempts: u32, session: Box, } -#[derive(Clone, Copy)] -struct RetryState { - revision: u64, - attempts: u32, -} - type Pending = Arc>>; -type FailedSnapshots = HashMap; type DeleteCallback = Box) + Send>; type PersistCallback = Box) + Send>; @@ -69,7 +63,6 @@ impl StorageWriter { .spawn(move || { let mut logs: HashMap = HashMap::new(); let mut durable_revisions: HashMap = HashMap::new(); - let mut retry_states: HashMap = HashMap::new(); loop { let op = if lock(&writer_pending).is_empty() { ops_rx.recv().ok() @@ -83,26 +76,20 @@ impl StorageWriter { let Some(op) = op else { break }; match op { Op::Flush => { - let failed = - flush(&writer_pending, &mut logs, &mut durable_revisions, &dir); - update_retry_states(&writer_pending, &mut retry_states, &failed); + flush(&writer_pending, &mut logs, &mut durable_revisions, &dir); } Op::Persist { session, done } => { - let failed = - flush(&writer_pending, &mut logs, &mut durable_revisions, &dir); - let result = persist_session( + flush(&writer_pending, &mut logs, &mut durable_revisions, &dir); + done(persist_session( &writer_pending, &mut logs, &mut durable_revisions, &dir, &session, - ); - update_retry_states(&writer_pending, &mut retry_states, &failed); - done(result); + )); } Op::Delete { id, done } => { lock(&writer_pending).remove(&id); - retry_states.remove(&id); logs.remove(&id); done(AppSession::delete(id, &dir)); } @@ -127,7 +114,14 @@ impl StorageWriter { .get(&session.id) .is_none_or(|current| current.revision <= revision); if replace { - pending.insert(session.id, PendingSnapshot { revision, session }); + pending.insert( + session.id, + PendingSnapshot { + revision, + retry_attempts: 0, + session, + }, + ); } drop(pending); if was_empty { @@ -181,50 +175,21 @@ fn writer_gone() -> SessionError { n00n_storage::StorageError::Io(io::Error::other("storage writer unavailable")).into() } -fn update_retry_states( - pending: &Pending, - retry_states: &mut HashMap, - failed: &FailedSnapshots, +fn retain_failed_snapshot( + pending: &mut HashMap, + id: N00nId, + mut snapshot: PendingSnapshot, ) { - let mut pending_guard = lock(pending); - retry_states.retain(|id, _| { - failed.get(id).is_some_and(|revision| { - pending_guard - .get(id) - .is_some_and(|snapshot| snapshot.revision == *revision) - }) - }); - let mut exhausted = Vec::new(); - for (&id, &revision) in failed { - if pending_guard - .get(&id) - .is_none_or(|snapshot| snapshot.revision != revision) - { - continue; - } - let state = retry_states.entry(id).or_insert(RetryState { - revision, - attempts: 0, - }); - if state.revision != revision { - state.revision = revision; - state.attempts = 0; - } - state.attempts += 1; - if state.attempts >= MAX_RETRY_ATTEMPTS { - exhausted.push((id, revision)); - } - } - - for (id, revision) in exhausted { - pending_guard.remove(&id); + snapshot.retry_attempts += 1; + if snapshot.retry_attempts >= MAX_RETRY_ATTEMPTS { warn!( - retry_count = MAX_RETRY_ATTEMPTS, + retry_count = snapshot.retry_attempts, %id, - revision, + revision = snapshot.revision, "storage writer exhausted retry attempts, dropping pending snapshot" ); - retry_states.remove(&id); + } else { + pending.insert(id, snapshot); } } @@ -233,46 +198,29 @@ fn flush( logs: &mut HashMap, durable_revisions: &mut HashMap, dir: &StateDir, -) -> FailedSnapshots { +) { let mut pending_guard = lock(pending); let batch = mem::take(&mut *pending_guard); if batch.is_empty() { - return HashMap::new(); + return; } let sessions_dir = match dir.ensure_subdir(SESSIONS_DIR) { Ok(d) => d, Err(e) => { warn!(error = %e, "failed to ensure sessions dir"); - let failed = batch - .iter() - .map(|(&id, snapshot)| (id, snapshot.revision)) - .collect(); for (id, snapshot) in batch { - let replace = pending_guard - .get(&id) - .is_none_or(|current| current.revision < snapshot.revision); - if replace { - pending_guard.insert(id, snapshot); - } + retain_failed_snapshot(&mut pending_guard, id, snapshot); } - return failed; + return; } }; - let mut failed = HashMap::new(); for (id, snapshot) in batch { if let Err(error) = write_session(&sessions_dir, logs, durable_revisions, &snapshot.session) { warn!(error = %error, id = %snapshot.session.id, "session write failed"); - let replace = pending_guard - .get(&id) - .is_none_or(|current| current.revision < snapshot.revision); - failed.insert(id, snapshot.revision); - if replace { - pending_guard.insert(id, snapshot); - } + retain_failed_snapshot(&mut pending_guard, id, snapshot); } } - failed } fn persist_session( @@ -489,20 +437,22 @@ mod tests { id, PendingSnapshot { revision, + retry_attempts: 0, session: Box::new(session), }, ); let mut logs = HashMap::new(); let mut durable_revisions = HashMap::new(); - assert!(!flush(&pending, &mut logs, &mut durable_revisions, &dir).is_empty()); - assert_eq!( - lock(&pending).get(&id).map(|item| item.revision), - Some(revision) - ); + flush(&pending, &mut logs, &mut durable_revisions, &dir); + let pending_guard = lock(&pending); + let snapshot = &pending_guard[&id]; + assert_eq!(snapshot.revision, revision); + assert_eq!(snapshot.retry_attempts, 1); + drop(pending_guard); std::fs::remove_dir(&blocked_log).unwrap(); - assert!(flush(&pending, &mut logs, &mut durable_revisions, &dir).is_empty()); + flush(&pending, &mut logs, &mut durable_revisions, &dir); assert!(lock(&pending).is_empty()); assert!(AppSession::load(id, &dir).is_ok()); } @@ -522,6 +472,7 @@ mod tests { id, PendingSnapshot { revision: newer.meta.revision, + retry_attempts: 0, session: Box::new(newer), }, ); @@ -542,69 +493,52 @@ mod tests { } #[test] - fn exhausted_retry_does_not_drop_newer_snapshot() { + fn failed_snapshot_is_dropped_after_retry_limit() { + let (_tmp, dir) = state_dir(); + let session = AppSession::new("test-model", "/tmp/exhausted"); + let id = session.id; + let sessions_dir = dir.ensure_subdir(SESSIONS_DIR).unwrap(); + std::fs::create_dir(sessions_dir.join(format!("{id}.jsonl"))).unwrap(); let pending: Pending = Arc::default(); - let mut older = AppSession::new("test-model", "/tmp/older"); - older.meta.revision = 1; - let id = older.id; - let mut newer = older.clone(); - newer.meta.revision = 2; lock(&pending).insert( id, PendingSnapshot { - revision: newer.meta.revision, - session: Box::new(newer), + revision: session.meta.revision, + retry_attempts: 0, + session: Box::new(session), }, ); + let mut logs = HashMap::new(); + let mut durable_revisions = HashMap::new(); - let failed = HashMap::from([(id, older.meta.revision)]); - let mut retry_states = HashMap::new(); for _ in 0..MAX_RETRY_ATTEMPTS { - update_retry_states(&pending, &mut retry_states, &failed); + flush(&pending, &mut logs, &mut durable_revisions, &dir); } - assert_eq!( - lock(&pending).get(&id).map(|snapshot| snapshot.revision), - Some(2) - ); - assert!(!retry_states.contains_key(&id)); + assert!(lock(&pending).is_empty()); } #[test] - fn retry_attempts_reset_for_newer_revision() { - let pending: Pending = Arc::default(); - let mut session = AppSession::new("test-model", "/tmp/revisions"); - session.meta.revision = 1; + fn equal_revision_replacement_resets_retry_attempts() { + let (_tmp, dir) = state_dir(); + let writer = StorageWriter::new(dir).unwrap(); + let session = AppSession::new("test-model", "/tmp/replacement"); let id = session.id; - let mut newer = session.clone(); - newer.meta.revision = 2; - lock(&pending).insert( - id, - PendingSnapshot { - revision: 1, - session: Box::new(session), - }, - ); - let mut retry_states = HashMap::new(); - let old_failure = HashMap::from([(id, 1)]); - for _ in 1..MAX_RETRY_ATTEMPTS { - update_retry_states(&pending, &mut retry_states, &old_failure); - } - lock(&pending).insert( + let revision = session.meta.revision; + lock(&writer.pending).insert( id, PendingSnapshot { - revision: 2, - session: Box::new(newer), + revision, + retry_attempts: MAX_RETRY_ATTEMPTS - 1, + session: Box::new(session.clone()), }, ); - let new_failure = HashMap::from([(id, 2)]); - update_retry_states(&pending, &mut retry_states, &new_failure); + writer.send(Box::new(session)); - let state = &retry_states[&id]; - assert_eq!(state.revision, 2); - assert_eq!(state.attempts, 1); - assert!(lock(&pending).contains_key(&id)); + assert_eq!(lock(&writer.pending)[&id].retry_attempts, 0); + lock(&writer.pending).clear(); + writer.shutdown(DRAIN_TIMEOUT); } #[test]