From e728c2dc3c756a22a27f4ba0d76adbe4cedec3fc Mon Sep 17 00:00:00 2001 From: Will Washburn Date: Sat, 19 Sep 2026 20:25:10 -0700 Subject: [PATCH 1/6] feat: report session capture and upload progress --- crates/ai-hist/src/ingest.rs | 173 +++++++++++++-- .../relayhistory/rust/src/probe/collector.rs | 53 +++-- plugins/relayhistory/rust/src/probe/main.rs | 10 +- .../relayhistory/rust/src/probe/progress.rs | 202 ++++++++++++++++++ 4 files changed, 400 insertions(+), 38 deletions(-) create mode 100644 plugins/relayhistory/rust/src/probe/progress.rs diff --git a/crates/ai-hist/src/ingest.rs b/crates/ai-hist/src/ingest.rs index 103fd38f..f2d10b45 100644 --- a/crates/ai-hist/src/ingest.rs +++ b/crates/ai-hist/src/ingest.rs @@ -84,6 +84,66 @@ pub fn sync_local_at(db_path: &Path) -> Result { sync_local_at_with_home(db_path, &home_dir()) } +/// Content-free progress for hosts displaying a local capture operation. +#[derive(Debug, Clone, Default, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct CaptureProgress { + pub source: String, + pub processed_files: usize, + pub total_files: Option, +} +thread_local! { + static CAPTURE_OBSERVER: std::cell::RefCell>> = std::cell::RefCell::new(None); +} + +/// Observes this thread's capture only; no paths or session contents are exposed. +pub fn sync_local_at_with_progress( + db_path: &Path, + observer: impl Fn(CaptureProgress) + 'static, +) -> Result { + struct Restore(Option>); + impl Drop for Restore { + fn drop(&mut self) { + CAPTURE_OBSERVER.with(|slot| *slot.borrow_mut() = self.0.take()); + } + } + let _restore = Restore(CAPTURE_OBSERVER.with(|slot| slot.replace(Some(Box::new(observer))))); + capture_progress("initializing", 0, None); + let result = sync_local_at(db_path); + if matches!(result, Ok(true)) { + capture_progress("complete", 0, None); + } + result +} + +fn capture_progress(source: &str, processed_files: usize, total_files: Option) { + CAPTURE_OBSERVER.with(|slot| { + if let Some(observer) = slot.borrow().as_ref() { + observer(CaptureProgress { + source: source.into(), + processed_files, + total_files, + }); + } + }); +} + +// Report before taking each file, including the final None. This counts completed +// files even when the ingest loop continues early for an unchanged checkpoint. +fn capture_files(source: &'static str, files: Vec) -> impl Iterator { + let total = files.len(); + let mut files = files.into_iter(); + let mut processed = 0; + std::iter::from_fn(move || { + capture_progress(source, processed, Some(total)); + let next = files.next(); + if next.is_some() { + processed += 1; + } + next + }) +} + /// Full local ingest using an explicit provider home instead of the process /// `HOME`. Used by [`crate::SessionStore`] when the embedder overrides home. pub(crate) fn sync_local_at_with_home(db_path: &Path, home: &Path) -> Result { @@ -419,6 +479,7 @@ fn sync_basic(conn: &Connection, db_path: &Path, home: &Path) -> Result<()> { // turns one interrupted run into a loop that re-scans from scratch forever // and never persists anything. Checkpointing makes each source's cursor // durable the moment that source completes. + capture_progress("claude-history", 0, None); if let Some(inserted) = report.capture( "claude", sync_jsonl_incremental( @@ -433,6 +494,7 @@ fn sync_basic(conn: &Connection, db_path: &Path, home: &Path) -> Result<()> { total_inserted += inserted; checkpoint_sync_state(&state_path, &state); } + capture_progress("claude", 0, None); if report .capture( "claude-metadata", @@ -442,10 +504,12 @@ fn sync_basic(conn: &Connection, db_path: &Path, home: &Path) -> Result<()> { { checkpoint_sync_state(&state_path, &state); } + capture_progress("codex", 0, None); if let Some(inserted) = report.capture("codex", sync_codex(conn, &mut state, home)) { total_inserted += inserted; checkpoint_sync_state(&state_path, &state); } + capture_progress("cursor", 0, None); if let Some(inserted) = report.capture( "cursor", sync_cursor(conn, &mut state, &home.join(".cursor/projects")), @@ -453,6 +517,7 @@ fn sync_basic(conn: &Connection, db_path: &Path, home: &Path) -> Result<()> { total_inserted += inserted; checkpoint_sync_state(&state_path, &state); } + capture_progress("grok", 0, None); if let Some(inserted) = report.capture( "grok", sync_grok(conn, &mut state, &home.join(".grok/sessions")), @@ -460,11 +525,13 @@ fn sync_basic(conn: &Connection, db_path: &Path, home: &Path) -> Result<()> { total_inserted += inserted; checkpoint_sync_state(&state_path, &state); } + capture_progress("trajectory", 0, None); if let Some(inserted) = report.capture("trajectory", sync_trajectories(conn, &mut state, home)) { total_inserted += inserted; checkpoint_sync_state(&state_path, &state); } + capture_progress("opencode", 0, None); let opencode = std::env::var_os("OPENCODE_DB") .map(PathBuf::from) .unwrap_or_else(|| home.join(".local/share/opencode/opencode.db")); @@ -480,6 +547,7 @@ fn sync_basic(conn: &Connection, db_path: &Path, home: &Path) -> Result<()> { // Establish connector-owned locators from actual provider enumeration after // ingestion, including on a checkpoint-only retry. Never infer an adapter // from an old aggregate presence row. + capture_progress("catalog", 0, None); let discovery_env = DiscoveryEnv::with_roots(conn, home.to_path_buf(), opencode); discover::discover_sessions_with_providers( &discovery_env, @@ -1571,7 +1639,7 @@ fn sync_codex_rollouts( if !root.exists() { continue; } - for rollout in collect_matching_files(&root, "rollout-", "jsonl")? { + for rollout in capture_files("codex", collect_matching_files(&root, "rollout-", "jsonl")?) { let key = rollout.to_string_lossy().to_string(); let stamp = file_stamp(&rollout)?; let record = seen.get(&key).and_then(Value::as_object); @@ -2651,7 +2719,7 @@ fn sync_claude_session_metadata( state.remove("claude_sessions_v2"); let mut scanned = 0; let mut upserted = 0; - for path in collect_matching_files(root, "", "jsonl")? { + for path in capture_files("claude", collect_matching_files(root, "", "jsonl")?) { let key = path.to_string_lossy().to_string(); let stamp = claude_sync_stamp(&path)?; if session_state.get(&key).and_then(Value::as_str) == Some(stamp.as_str()) @@ -3972,7 +4040,10 @@ fn sync_grok(conn: &Connection, state: &mut Map, root: &Path) -> let mut scanned = 0; let mut sessions = 0; let mut errors = 0; - for chat in collect_matching_files(root, "chat_history", "jsonl")? { + for chat in capture_files( + "grok", + collect_matching_files(root, "chat_history", "jsonl")?, + ) { let key = chat.to_string_lossy().to_string(); let stamp = grok_session_stamp(&chat)?; if grok_state.get(&key).and_then(Value::as_str) == Some(stamp.as_str()) { @@ -4227,7 +4298,7 @@ fn sync_trajectories( let mut updated = 0; let mut skipped = 0; let mut errors = 0; - for path in files { + for path in capture_files("trajectory", files) { let metadata = match path.metadata() { Ok(metadata) => metadata, Err(_) => { @@ -4344,11 +4415,27 @@ fn collect_named_dirs(root: &Path, name: &str, out: &mut Vec) -> Result return Ok(()); } for entry in fs::read_dir(root)? { - let path = entry?.path(); - if path.is_dir() { - if path.file_name().and_then(|s| s.to_str()) == Some(name) { - out.push(path.clone()); - } + let entry = entry?; + // Never follow symlinks: dependency links can revisit the same tree or cycle. + if !entry.file_type()?.is_dir() { + continue; + } + let path = entry.path(); + let file_name = entry.file_name(); + let child = file_name.to_str().unwrap_or(""); + if child == name { + out.push(path); + } else if !matches!( + child, + "node_modules" + | ".git" + | "target" + | ".next" + | ".venv" + | "venv" + | "__pycache__" + | ".cache" + ) { collect_named_dirs(&path, name, out)?; } } @@ -4363,10 +4450,12 @@ fn collect_trajectory_json(dir: &Path, out: &mut Vec) -> Result<()> { return Ok(()); } for entry in fs::read_dir(dir)? { - let path = entry?.path(); - if path.is_dir() { + let entry = entry?; + let file_type = entry.file_type()?; + let path = entry.path(); + if file_type.is_dir() { collect_trajectory_json(&path, out)?; - } else if path.extension().and_then(|s| s.to_str()) == Some("json") { + } else if file_type.is_file() && path.extension().and_then(|s| s.to_str()) == Some("json") { let name = path.file_name().and_then(|s| s.to_str()).unwrap_or(""); if name != "index.json" && name != ".sync-state.json" && !name.ends_with(".trace.json") { @@ -8094,3 +8183,63 @@ mod tests { ); } } + +#[cfg(test)] +mod capture_progress_tests { + use super::*; + use std::sync::{Arc, Mutex}; + + #[test] + fn file_progress_counts_completed_files_and_restores_observer() { + let updates = Arc::new(Mutex::new(Vec::new())); + let observed = updates.clone(); + CAPTURE_OBSERVER.with(|slot| { + *slot.borrow_mut() = Some(Box::new(move |p| observed.lock().unwrap().push(p))) + }); + let files: Vec<_> = capture_files("claude", vec!["a".into(), "b".into()]).collect(); + assert_eq!(files.len(), 2); + let values = updates.lock().unwrap(); + assert_eq!( + values.iter().map(|v| v.processed_files).collect::>(), + vec![0, 1, 2] + ); + assert!(values.iter().all(|v| v.total_files == Some(2))); + drop(values); + // Even an early database failure restores the caller's observer. + let home = tempfile::tempdir().unwrap(); + let db = home.path().join("bad.db"); + fs::write(&db, "not a database").unwrap(); + assert!(sync_local_at_with_progress(&db, |_| {}).is_err()); + capture_progress("restored", 0, None); + assert_eq!(updates.lock().unwrap().last().unwrap().source, "restored"); + CAPTURE_OBSERVER.with(|slot| slot.borrow_mut().take()); + } + + #[test] + fn trajectory_discovery_prunes_dependencies_and_build_output() { + let home = tempfile::tempdir().unwrap(); + let wanted = home.path().join("repo/.trajectories"); + fs::create_dir_all(wanted.join("completed/month")).unwrap(); + for ignored in [ + "node_modules/pkg", + ".git/objects", + "target/debug", + ".next/cache", + ".venv/lib", + ] { + fs::create_dir_all(home.path().join(ignored).join(".trajectories")).unwrap(); + } + #[cfg(unix)] + { + std::os::unix::fs::symlink(home.path(), home.path().join("repo/cycle")).unwrap(); + std::os::unix::fs::symlink(&wanted, wanted.join("completed/cycle")).unwrap(); + } + let mut roots = vec![]; + collect_named_dirs(home.path(), ".trajectories", &mut roots).unwrap(); + assert_eq!(roots, vec![wanted.clone()]); + fs::write(wanted.join("completed/month/run.json"), "{}").unwrap(); + let mut files = vec![]; + collect_trajectory_json(&wanted, &mut files).unwrap(); + assert_eq!(files, vec![wanted.join("completed/month/run.json")]); + } +} diff --git a/plugins/relayhistory/rust/src/probe/collector.rs b/plugins/relayhistory/rust/src/probe/collector.rs index 8ab0eb83..621e3289 100644 --- a/plugins/relayhistory/rust/src/probe/collector.rs +++ b/plugins/relayhistory/rust/src/probe/collector.rs @@ -1,7 +1,7 @@ use super::{lock, read_config, save_json, user_error, Config}; use ai_hist::delivery::{self, worker, ExportSelection}; use anyhow::{ensure, Context, Result}; -use relayhistory_plugin::{cloud, destination}; +use relayhistory_plugin::destination; use rusqlite::Connection; use serde_json::json; use std::{ @@ -57,8 +57,7 @@ pub fn record_baseline(conn: &Connection, include_existing: bool) -> Result Result } Ok(status) } +pub fn capture(directory: &Path, history_url: &str) -> Result<()> { + let progress = super::progress::Monitor::start(directory, history_url, None); + let result = + ai_hist::sync_local_at_with_progress(&directory.join("history.db"), progress.observer()); + progress.finish(matches!(result, Ok(true))); + ensure!(result?, "capture not complete"); + Ok(()) +} + pub fn cycle(directory: &Path, config: &Config) -> Result<()> { ensure!( destination::selected_account(Some(&config.history_url))? == config.delivery_account, "wrong destination" ); - let db_path = directory.join("history.db"); + capture(directory, &config.history_url)?; + deliver_captured(directory, config) +} + +pub fn deliver_captured(directory: &Path, config: &Config) -> Result<()> { ensure!( - ai_hist::sync_local_at(&db_path)?, - "capture not complete" + destination::selected_account(Some(&config.history_url))? == config.delivery_account, + "wrong destination" ); + let db_path = directory.join("history.db"); // The receiver rejects a batch whose account or mapping does not match, but // only once one exists. An idle generation pointed at another destination // must not look healthy, so the saved configuration is checked outright. @@ -149,17 +162,11 @@ pub fn cycle(directory: &Path, config: &Config) -> Result<()> { // its verdict is a generic permission refusal. Name the cause here first so // the user sees which uploader to stop instead of a reconnect suggestion. super::check_legacy_schedules(config.acknowledge_uninspected_schedules)?; - let status = deliver(&db_path, config)?; - let token = cloud::access_token(Some(&config.history_url))?; - let agent = ureq::AgentBuilder::new() - .redirects(0) - .timeout(Duration::from_secs(15)) - .build(); - agent - .post(&format!("{}/v1/onboarding/heartbeat", config.history_url)) - .set("Authorization", &format!("Bearer {token}")) - .call() - .map_err(|_| user_error("Could not confirm the connection with Cloud."))?; + let progress = + super::progress::Monitor::start(directory, &config.history_url, Some(&config.job_id)); + let result = deliver(&db_path, config); + progress.finish(result.is_ok()); + let status = result?; println!( "Probe connected: {} records received, {} queued.", status.acknowledged_records, status.pending_records @@ -250,8 +257,15 @@ pub fn run_background(directory: &Path, startup_id: &str) -> Result<()> { &directory.join("runtime.json"), &json!({"startup_id":startup_id,"pid":std::process::id(),"ready":true}), )?; + let mut capture_due = Instant::now(); while !stop_requested(directory, startup_id) { - if cycle(directory, &config).is_err() { + let result = if Instant::now() >= capture_due { + capture_due = Instant::now() + Duration::from_secs(60); + cycle(directory, &config) + } else { + deliver_captured(directory, &config) + }; + if result.is_err() { eprintln!("Sync paused or offline. Retrying; local data remains queued."); } for _ in 0..20 { @@ -380,8 +394,7 @@ mod tests { ) .unwrap(); } - let baseline = - ai_hist::storage::session_identities_after(&conn, None, 10_000).unwrap(); + let baseline = ai_hist::storage::session_identities_after(&conn, None, 10_000).unwrap(); // The same identities carried inline would exceed the 64 KiB cap that // create_job enforces on a delivery configuration. assert!(serde_json::to_vec(&baseline).unwrap().len() > 65_536); diff --git a/plugins/relayhistory/rust/src/probe/main.rs b/plugins/relayhistory/rust/src/probe/main.rs index 2243322d..7dcf80b0 100644 --- a/plugins/relayhistory/rust/src/probe/main.rs +++ b/plugins/relayhistory/rust/src/probe/main.rs @@ -1,5 +1,6 @@ //! Standalone Cloud collector. Reuses ai-hist capture/queue and the optional transport. mod collector; +mod progress; use anyhow::{ensure, Context, Result}; use clap::{Args, Parser, Subcommand}; @@ -260,7 +261,7 @@ fn install(options: Install) -> Result<()> { cloud::CloudBearerOptions { force_login: options.force_login, interactive: true, - client_name: "Agent Relay Probe", + client_name: "Agent Relay Session Recorder", announce: &mut |approval: &cloud::DeviceApproval| { println!( "Open this URL to authorize your computer:\n{}", @@ -340,10 +341,7 @@ fn install(options: Install) -> Result<()> { cloud::save_auth(&session)?; let db_path = directory.join("history.db"); println!("Preparing local session capture…"); - ensure!( - ai_hist::sync_local_at(&db_path)?, - "capture did not complete" - ); + collector::capture(&directory, &history_url)?; let conn = ai_hist::open_db(&db_path)?; let config = match existing { Some(mut config) => { @@ -416,7 +414,7 @@ fn install(options: Install) -> Result<()> { } }; drop(conn); - collector::cycle(&directory, &config)?; + collector::deliver_captured(&directory, &config)?; if options.once { println!("One capture/delivery cycle completed."); return Ok(()); diff --git a/plugins/relayhistory/rust/src/probe/progress.rs b/plugins/relayhistory/rust/src/probe/progress.rs new file mode 100644 index 00000000..2c7f1327 --- /dev/null +++ b/plugins/relayhistory/rust/src/probe/progress.rs @@ -0,0 +1,202 @@ +//! Content-free, best-effort progress. A slow/offline status endpoint never blocks capture. +use ai_hist::delivery; +use relayhistory_plugin::cloud; +use serde::Serialize; +use std::{ + path::Path, + sync::{mpsc, Arc, Mutex}, + thread, + time::{Duration, Instant}, +}; + +#[derive(Clone, Default, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct Progress { + pub phase: String, + pub source: String, + pub processed_files: usize, + pub total_files: Option, + pub sessions_captured: i64, + pub records_uploaded: i64, + pub records_queued: i64, + pub backlog_complete: bool, +} + +impl Progress { + fn read_counts(&mut self, db: &Path, job: Option<&str>) { + let Ok(conn) = ai_hist::open_db_readonly(db) else { + return; + }; + let _ = conn.busy_timeout(Duration::from_millis(100)); + self.sessions_captured = conn + .query_row( + "SELECT COUNT(*) FROM sessions WHERE discovery_state = 'full'", + [], + |r| r.get(0), + ) + .unwrap_or(0); + if let Some(job) = job { + if let Ok(status) = delivery::status(&conn, job) { + self.apply_delivery(&status); + } + } + } + fn apply_delivery(&mut self, status: &delivery::DeliveryStatus) { + self.records_uploaded = status.acknowledged_records; + self.records_queued = status.pending_records; + self.backlog_complete = status.bootstrap_complete && status.unqueued_changes == 0; + self.phase = if status.state != "active" || status.failure.is_some() { + "paused" + } else if self.backlog_complete && status.pending_records == 0 { + "watching" + } else { + "uploading" + } + .into(); + } + fn line(&self) -> String { + match self.phase.as_str() { + "scanning" => match self.total_files { + Some(total) => format!( + "Reading {} session files: {} of {} processed · {} sessions captured", + self.source, self.processed_files, total, self.sessions_captured + ), + None => format!( + "Scanning {}… {} sessions captured", + self.source, self.sessions_captured + ), + }, + "paused" => format!( + "Upload paused; retrying. {} records uploaded, {} queued", + self.records_uploaded, self.records_queued + ), + "watching" => format!( + "Up to date: {} records uploaded. Watching for new sessions.", + self.records_uploaded + ), + _ => format!( + "Uploading: {} records received, {} queued{}", + self.records_uploaded, + self.records_queued, + if self.backlog_complete { + "" + } else { + " · preparing remaining records" + } + ), + } + } +} + +pub struct Monitor { + snapshot: Arc>, + stop: Option>, + thread: Option>, +} +impl Monitor { + pub fn start(directory: &Path, history_url: &str, job: Option<&str>) -> Self { + let snapshot = Arc::new(Mutex::new(Progress { + phase: if job.is_some() { + "uploading" + } else { + "scanning" + } + .into(), + source: "local history".into(), + ..Default::default() + })); + let shared = snapshot.clone(); + let db = directory.join("history.db"); + let url = history_url.to_owned(); + let job = job.map(str::to_owned); + let (stop, receive) = mpsc::channel(); + let thread = thread::spawn(move || { + let started = Instant::now(); + let mut finish = None; + loop { + let mut progress = shared.lock().unwrap().clone(); + progress.read_counts(&db, job.as_deref()); + if finish == Some(false) { + progress.phase = "paused".into(); + } + println!("{} ({}s)", progress.line(), started.elapsed().as_secs()); + heartbeat(&url, &progress); + if finish.is_some() { + break; + } + match receive.recv_timeout(Duration::from_secs(3)) { + Ok(success) => finish = Some(success), + Err(mpsc::RecvTimeoutError::Disconnected) => finish = Some(false), + Err(mpsc::RecvTimeoutError::Timeout) => {} + } + } + }); + Self { + snapshot, + stop: Some(stop), + thread: Some(thread), + } + } + pub fn observer(&self) -> impl Fn(ai_hist::CaptureProgress) + 'static { + let snapshot = self.snapshot.clone(); + move |capture| { + let mut p = snapshot.lock().unwrap(); + p.source = capture.source; + p.processed_files = capture.processed_files; + p.total_files = capture.total_files; + } + } + pub fn finish(mut self, success: bool) { + if let Some(stop) = self.stop.take() { + let _ = stop.send(success); + } + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} +impl Drop for Monitor { + fn drop(&mut self) { + self.stop.take(); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} +fn heartbeat(url: &str, progress: &Progress) { + let Ok(token) = cloud::access_token(Some(url)) else { + return; + }; + let result = ureq::AgentBuilder::new() + .redirects(0) + .timeout(Duration::from_secs(3)) + .build() + .post(&format!("{url}/v1/onboarding/heartbeat")) + .set("Authorization", &format!("Bearer {token}")) + .send_json(serde_json::json!({"progress":progress})); + if result.is_err() { + eprintln!("Cloud progress update unavailable; collection continues locally."); + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn unknown_totals_do_not_claim_completion() { + let progress = Progress { + phase: "uploading".into(), + records_uploaded: 20, + ..Default::default() + }; + assert!(progress.line().contains("preparing remaining records")); + let capture = Progress { + phase: "scanning".into(), + source: "claude".into(), + processed_files: 2, + total_files: Some(5), + ..Default::default() + }; + assert!(capture.line().contains("2 of 5")); + } +} From a3d8e3fedd02c80b615ca25725f01570af2427dd Mon Sep 17 00:00:00 2001 From: Will Washburn Date: Sat, 19 Sep 2026 22:09:57 -0700 Subject: [PATCH 2/6] Make capture progress reentrant and heartbeat shutdown nonblocking --- crates/ai-hist/src/ingest.rs | 43 ++++++- .../relayhistory/rust/src/probe/progress.rs | 112 +++++++++++++++--- 2 files changed, 136 insertions(+), 19 deletions(-) diff --git a/crates/ai-hist/src/ingest.rs b/crates/ai-hist/src/ingest.rs index f2d10b45..58908f70 100644 --- a/crates/ai-hist/src/ingest.rs +++ b/crates/ai-hist/src/ingest.rs @@ -92,8 +92,9 @@ pub struct CaptureProgress { pub processed_files: usize, pub total_files: Option, } +type CaptureObserver = std::rc::Rc; thread_local! { - static CAPTURE_OBSERVER: std::cell::RefCell>> = std::cell::RefCell::new(None); + static CAPTURE_OBSERVER: std::cell::RefCell> = std::cell::RefCell::new(None); } /// Observes this thread's capture only; no paths or session contents are exposed. @@ -101,13 +102,14 @@ pub fn sync_local_at_with_progress( db_path: &Path, observer: impl Fn(CaptureProgress) + 'static, ) -> Result { - struct Restore(Option>); + struct Restore(Option); impl Drop for Restore { fn drop(&mut self) { CAPTURE_OBSERVER.with(|slot| *slot.borrow_mut() = self.0.take()); } } - let _restore = Restore(CAPTURE_OBSERVER.with(|slot| slot.replace(Some(Box::new(observer))))); + let _restore = + Restore(CAPTURE_OBSERVER.with(|slot| slot.replace(Some(std::rc::Rc::new(observer))))); capture_progress("initializing", 0, None); let result = sync_local_at(db_path); if matches!(result, Ok(true)) { @@ -118,7 +120,8 @@ pub fn sync_local_at_with_progress( fn capture_progress(source: &str, processed_files: usize, total_files: Option) { CAPTURE_OBSERVER.with(|slot| { - if let Some(observer) = slot.borrow().as_ref() { + let observer = slot.borrow().clone(); + if let Some(observer) = observer { observer(CaptureProgress { source: source.into(), processed_files, @@ -8194,7 +8197,7 @@ mod capture_progress_tests { let updates = Arc::new(Mutex::new(Vec::new())); let observed = updates.clone(); CAPTURE_OBSERVER.with(|slot| { - *slot.borrow_mut() = Some(Box::new(move |p| observed.lock().unwrap().push(p))) + *slot.borrow_mut() = Some(std::rc::Rc::new(move |p| observed.lock().unwrap().push(p))) }); let files: Vec<_> = capture_files("claude", vec!["a".into(), "b".into()]).collect(); assert_eq!(files.len(), 2); @@ -8215,6 +8218,36 @@ mod capture_progress_tests { CAPTURE_OBSERVER.with(|slot| slot.borrow_mut().take()); } + #[test] + fn callbacks_can_reenter_capture_and_restore_the_outer_observer() { + let home = tempfile::tempdir().unwrap(); + let db = home.path().join("bad.db"); + fs::write(&db, "not a database").unwrap(); + let inner_db = db.clone(); + let updates = Arc::new(Mutex::new(Vec::new())); + let outer_updates = updates.clone(); + assert!(sync_local_at_with_progress(&db, move |progress| { + outer_updates + .lock() + .unwrap() + .push(format!("outer:{}", progress.source)); + let inner_updates = outer_updates.clone(); + assert!(sync_local_at_with_progress(&inner_db, move |inner| { + inner_updates + .lock() + .unwrap() + .push(format!("inner:{}", inner.source)); + }) + .is_err()); + }) + .is_err()); + assert_eq!( + *updates.lock().unwrap(), + vec!["outer:initializing", "inner:initializing"] + ); + CAPTURE_OBSERVER.with(|slot| assert!(slot.borrow().is_none())); + } + #[test] fn trajectory_discovery_prunes_dependencies_and_build_output() { let home = tempfile::tempdir().unwrap(); diff --git a/plugins/relayhistory/rust/src/probe/progress.rs b/plugins/relayhistory/rust/src/probe/progress.rs index 2c7f1327..3b34f2cf 100644 --- a/plugins/relayhistory/rust/src/probe/progress.rs +++ b/plugins/relayhistory/rust/src/probe/progress.rs @@ -30,7 +30,7 @@ impl Progress { let _ = conn.busy_timeout(Duration::from_millis(100)); self.sessions_captured = conn .query_row( - "SELECT COUNT(*) FROM sessions WHERE discovery_state = 'full'", + "SELECT COUNT(*) FROM sessions WHERE COALESCE(discovery_state, 'full') = 'full'", [], |r| r.get(0), ) @@ -66,6 +66,10 @@ impl Progress { self.source, self.sessions_captured ), }, + "capture_paused" => format!( + "Local capture paused; retrying. {} session files processed, {} sessions captured", + self.processed_files, self.sessions_captured + ), "paused" => format!( "Upload paused; retrying. {} records uploaded, {} queued", self.records_uploaded, self.records_queued @@ -91,10 +95,17 @@ impl Progress { pub struct Monitor { snapshot: Arc>, stop: Option>, - thread: Option>, } impl Monitor { pub fn start(directory: &Path, history_url: &str, job: Option<&str>) -> Self { + let url = history_url.to_owned(); + Self::start_with_report(directory, job, move |progress| heartbeat(&url, progress)) + } + fn start_with_report( + directory: &Path, + job: Option<&str>, + mut report: impl FnMut(&Progress) + Send + 'static, + ) -> Self { let snapshot = Arc::new(Mutex::new(Progress { phase: if job.is_some() { "uploading" @@ -107,20 +118,29 @@ impl Monitor { })); let shared = snapshot.clone(); let db = directory.join("history.db"); - let url = history_url.to_owned(); let job = job.map(str::to_owned); let (stop, receive) = mpsc::channel(); - let thread = thread::spawn(move || { + thread::spawn(move || { let started = Instant::now(); let mut finish = None; loop { + // The next upload monitor now owns progress. Do not send a + // delayed final "scanning" heartbeat after successful capture. + if finish == Some(true) && job.is_none() { + break; + } let mut progress = shared.lock().unwrap().clone(); progress.read_counts(&db, job.as_deref()); if finish == Some(false) { - progress.phase = "paused".into(); + progress.phase = if job.is_none() { + "capture_paused" + } else { + "paused" + } + .into(); } println!("{} ({}s)", progress.line(), started.elapsed().as_secs()); - heartbeat(&url, &progress); + report(&progress); if finish.is_some() { break; } @@ -134,7 +154,6 @@ impl Monitor { Self { snapshot, stop: Some(stop), - thread: Some(thread), } } pub fn observer(&self) -> impl Fn(ai_hist::CaptureProgress) + 'static { @@ -146,27 +165,40 @@ impl Monitor { p.total_files = capture.total_files; } } + // Sending stop is sufficient: the detached worker exits after its bounded + // request/final report. Capture, delivery and shutdown never join network I/O. pub fn finish(mut self, success: bool) { if let Some(stop) = self.stop.take() { let _ = stop.send(success); } - if let Some(thread) = self.thread.take() { - let _ = thread.join(); - } } } impl Drop for Monitor { fn drop(&mut self) { self.stop.take(); - if let Some(thread) = self.thread.take() { - let _ = thread.join(); - } } } fn heartbeat(url: &str, progress: &Progress) { - let Ok(token) = cloud::access_token(Some(url)) else { + // Progress must never acquire the refresh lock or rotate credentials. The + // delivery path owns renewal; skip this heartbeat if its cached token expired. + let Ok(Some(auth)) = cloud::load_selected_auth(Some(url)) else { return; }; + let valid = auth + .access_token_expires_at + .as_deref() + .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) + .is_some_and(|expiry| expiry.timestamp() > chrono::Utc::now().timestamp() + 5); + if !valid + || !auth.access_token.starts_with("rth_at_") + || !auth + .access_token + .bytes() + .all(|byte| byte.is_ascii_graphic()) + { + return; + } + let token = auth.access_token; let result = ureq::AgentBuilder::new() .redirects(0) .timeout(Duration::from_secs(3)) @@ -182,6 +214,58 @@ fn heartbeat(url: &str, progress: &Progress) { #[cfg(test)] mod tests { use super::*; + #[test] + fn finish_and_drop_do_not_wait_for_a_slow_reporter() { + for success in [None, Some(false), Some(true)] { + let directory = tempfile::tempdir().unwrap(); + let (entered, started) = mpsc::channel(); + let (release, blocked) = mpsc::channel(); + let (finished, done) = mpsc::channel(); + let mut first = true; + let monitor = Monitor::start_with_report(directory.path(), None, move |progress| { + if first { + first = false; + entered.send(()).unwrap(); + blocked.recv().unwrap(); + } else { + finished.send(progress.phase.clone()).unwrap(); + } + }); + started.recv_timeout(Duration::from_secs(2)).unwrap(); + let start = Instant::now(); + if let Some(success) = success { + monitor.finish(success); + } else { + drop(monitor); + } + assert!(start.elapsed() < Duration::from_millis(100)); + release.send(()).unwrap(); + let final_phase = done.recv_timeout(Duration::from_secs(2)); + if success == Some(true) { + assert!(matches!( + final_phase, + Err(mpsc::RecvTimeoutError::Disconnected) + )); + } else { + assert_eq!(final_phase.unwrap(), "capture_paused"); + } + } + } + + #[test] + fn legacy_null_discovery_state_counts_as_captured() { + let directory = tempfile::tempdir().unwrap(); + let db = directory.path().join("history.db"); + let conn = rusqlite::Connection::open(&db).unwrap(); + conn.execute_batch("CREATE TABLE sessions (discovery_state TEXT); INSERT INTO sessions VALUES (NULL), ('full'), ('partial');").unwrap(); + let mut progress = Progress::default(); + progress.read_counts(&db, None); + assert_eq!(progress.sessions_captured, 2); + progress.phase = "capture_paused".into(); + assert!(progress.line().contains("Local capture paused")); + assert!(!progress.line().contains("uploaded")); + } + #[test] fn unknown_totals_do_not_claim_completion() { let progress = Progress { From c8ff52b26c73543f4664812a07eb84399ff96944 Mon Sep 17 00:00:00 2001 From: Will Washburn Date: Sat, 19 Sep 2026 22:17:53 -0700 Subject: [PATCH 3/6] Keep idle heartbeats alive and flush one-shot progress with a deadline --- plugins/relayhistory/rust/src/cloud.rs | 139 ++++++++++++++++-- .../relayhistory/rust/src/probe/collector.rs | 12 +- plugins/relayhistory/rust/src/probe/main.rs | 2 +- .../relayhistory/rust/src/probe/progress.rs | 97 ++++++++---- 4 files changed, 209 insertions(+), 41 deletions(-) diff --git a/plugins/relayhistory/rust/src/cloud.rs b/plugins/relayhistory/rust/src/cloud.rs index e50db5ec..7077c7d7 100644 --- a/plugins/relayhistory/rust/src/cloud.rs +++ b/plugins/relayhistory/rust/src/cloud.rs @@ -359,6 +359,51 @@ pub fn access_token(base_url: Option<&str>) -> Result { Ok(auth.access_token) } +/// Best-effort renewal for idle progress reporting. Never wait for another +/// transport's refresh lock, and bound the refresh request including its body. +/// The rotated credential pair is persisted through the same atomic store. +pub fn try_progress_access_token( + base_url: &str, + timeout: std::time::Duration, +) -> Result> { + let Some(mut auth) = load_selected_auth(Some(base_url))? else { + return Ok(None); + }; + if !token_has_valid_lifetime(&auth) { + let path = refresh_lock_path(base_url)?; + let file = fs::OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(path)?; + if let Err(error) = fs2::FileExt::try_lock_exclusive(&file) { + if error.kind() == std::io::ErrorKind::WouldBlock { + return Ok(None); + } + return Err(error.into()); + } + let _lock = AuthRefreshLock { _file: file }; + auth = load_selected_auth(Some(base_url))?.context("progress credentials disappeared")?; + if !token_has_valid_lifetime(&auth) { + let agent = ureq::AgentBuilder::new() + .redirects(0) + .timeout(timeout) + .build(); + auth = refresh_auth_with_agent(&auth, &agent) + .map_err(|_| anyhow::anyhow!("progress credential renewal unavailable"))?; + save_auth(&auth)?; + } + } + Ok((token_has_valid_lifetime(&auth) + && auth.access_token.starts_with("rth_at_") + && auth + .access_token + .bytes() + .all(|byte| byte.is_ascii_graphic())) + .then_some(auth.access_token)) +} + fn token_has_valid_lifetime(auth: &StoredAuth) -> bool { let expiry = auth .access_token_expires_at @@ -445,10 +490,7 @@ pub fn machine_id() -> Result { .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_nanos()) .unwrap_or(0); - let id = format!( - "m_{}", - ai_hist::prompt_hash(&format!("{host}:{nanos}")) - ); + let id = format!("m_{}", ai_hist::prompt_hash(&format!("{host}:{nanos}"))); write_private(&path, &id)?; Ok(id) } @@ -985,6 +1027,10 @@ fn refresh_and_save_auth(auth: &StoredAuth) -> Result { } fn refresh_auth(auth: &StoredAuth) -> Result { + refresh_auth_with_agent(auth, &ureq::AgentBuilder::new().build()) +} + +fn refresh_auth_with_agent(auth: &StoredAuth, agent: &ureq::Agent) -> Result { require_secure_transport(&auth.base_url)?; let refresh_token = auth .refresh_token @@ -994,7 +1040,8 @@ fn refresh_auth(auth: &StoredAuth) -> Result { "{}/v1/auth/token/refresh", auth.base_url.trim_end_matches('/') ); - let response = ureq::post(&url) + let response = agent + .post(&url) .set("Content-Type", "application/json") .send_json(serde_json::json!({ "refreshToken": refresh_token })) .map_err(map_http_err)?; @@ -1684,10 +1731,7 @@ pub fn resolve_recall_auth( auth.access_token.starts_with("rth_at_"), "stored relayhistory session has no rth_at_ access token (run `ai-hist login`)" ); - let needs_tenancy = auth - .org_id - .as_deref() - .is_none_or(|id| id.trim().is_empty()); + let needs_tenancy = auth.org_id.as_deref().is_none_or(|id| id.trim().is_empty()); // A session created before the service began returning tenancy can still // recover it on refresh. Do this before the provenance check, under the // same lock as every other token rotation, so a concurrent writer cannot @@ -3285,6 +3329,83 @@ pub(crate) mod tests { }); } + #[test] + fn idle_progress_refresh_is_bounded_and_respects_the_rotation_lock() { + with_temp_home(|| { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + save_auth(&StoredAuth { + base_url: base_url.clone(), + access_token: "rth_at_expired".into(), + refresh_token: Some("rth_rt_old".into()), + ..Default::default() + }) + .unwrap(); + let lock = acquire_refresh_lock(&base_url).unwrap(); + let start = std::time::Instant::now(); + assert!( + try_progress_access_token(&base_url, std::time::Duration::from_millis(100)) + .unwrap() + .is_none() + ); + assert!(start.elapsed() < std::time::Duration::from_millis(100)); + drop(lock); + // The listening socket never responds: the refresh has a hard deadline. + let start = std::time::Instant::now(); + assert!( + try_progress_access_token(&base_url, std::time::Duration::from_millis(100)) + .is_err() + ); + assert!(start.elapsed() < std::time::Duration::from_secs(1)); + }); + } + + #[test] + fn idle_progress_refresh_persists_and_reuses_the_rotated_pair() { + with_temp_home(|| { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + save_auth(&StoredAuth { + base_url: base_url.clone(), + access_token: "rth_at_expired".into(), + refresh_token: Some("rth_rt_old".into()), + ..Default::default() + }) + .unwrap(); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let (line, _, _) = read_http_request(&mut stream); + assert!(line.starts_with("POST /v1/auth/token/refresh ")); + write_http_response(&mut stream, "200 OK", &serde_json::json!({ + "accessToken": "rth_at_fresh", "refreshToken": "rth_rt_fresh", + "accessTokenExpiresAt": (chrono::Utc::now() + chrono::Duration::hours(1)).to_rfc3339() + }).to_string()); + }); + assert_eq!( + try_progress_access_token(&base_url, std::time::Duration::from_secs(1)) + .unwrap() + .as_deref(), + Some("rth_at_fresh") + ); + server.join().unwrap(); + // The listener is gone; a valid cached credential needs no network call. + assert_eq!( + try_progress_access_token(&base_url, std::time::Duration::from_secs(1)) + .unwrap() + .as_deref(), + Some("rth_at_fresh") + ); + assert_eq!( + load_selected_auth(Some(&base_url)) + .unwrap() + .unwrap() + .refresh_token + .as_deref(), + Some("rth_rt_fresh") + ); + }); + } + #[test] fn a_reused_auth_value_adopts_the_persisted_rotation() { with_temp_home(|| { diff --git a/plugins/relayhistory/rust/src/probe/collector.rs b/plugins/relayhistory/rust/src/probe/collector.rs index 621e3289..4050e719 100644 --- a/plugins/relayhistory/rust/src/probe/collector.rs +++ b/plugins/relayhistory/rust/src/probe/collector.rs @@ -137,10 +137,10 @@ pub fn cycle(directory: &Path, config: &Config) -> Result<()> { "wrong destination" ); capture(directory, &config.history_url)?; - deliver_captured(directory, config) + deliver_captured(directory, config, false) } -pub fn deliver_captured(directory: &Path, config: &Config) -> Result<()> { +pub fn deliver_captured(directory: &Path, config: &Config, before_exit: bool) -> Result<()> { ensure!( destination::selected_account(Some(&config.history_url))? == config.delivery_account, "wrong destination" @@ -165,7 +165,11 @@ pub fn deliver_captured(directory: &Path, config: &Config) -> Result<()> { let progress = super::progress::Monitor::start(directory, &config.history_url, Some(&config.job_id)); let result = deliver(&db_path, config); - progress.finish(result.is_ok()); + if before_exit { + progress.finish_before_exit(result.is_ok(), Duration::from_secs(5)); + } else { + progress.finish(result.is_ok()); + } let status = result?; println!( "Probe connected: {} records received, {} queued.", @@ -263,7 +267,7 @@ pub fn run_background(directory: &Path, startup_id: &str) -> Result<()> { capture_due = Instant::now() + Duration::from_secs(60); cycle(directory, &config) } else { - deliver_captured(directory, &config) + deliver_captured(directory, &config, false) }; if result.is_err() { eprintln!("Sync paused or offline. Retrying; local data remains queued."); diff --git a/plugins/relayhistory/rust/src/probe/main.rs b/plugins/relayhistory/rust/src/probe/main.rs index 7dcf80b0..e8e26348 100644 --- a/plugins/relayhistory/rust/src/probe/main.rs +++ b/plugins/relayhistory/rust/src/probe/main.rs @@ -414,7 +414,7 @@ fn install(options: Install) -> Result<()> { } }; drop(conn); - collector::deliver_captured(&directory, &config)?; + collector::deliver_captured(&directory, &config, true)?; if options.once { println!("One capture/delivery cycle completed."); return Ok(()); diff --git a/plugins/relayhistory/rust/src/probe/progress.rs b/plugins/relayhistory/rust/src/probe/progress.rs index 3b34f2cf..0fe97be1 100644 --- a/plugins/relayhistory/rust/src/probe/progress.rs +++ b/plugins/relayhistory/rust/src/probe/progress.rs @@ -95,16 +95,19 @@ impl Progress { pub struct Monitor { snapshot: Arc>, stop: Option>, + completion: mpsc::Receiver<()>, } impl Monitor { pub fn start(directory: &Path, history_url: &str, job: Option<&str>) -> Self { let url = history_url.to_owned(); - Self::start_with_report(directory, job, move |progress| heartbeat(&url, progress)) + Self::start_with_report(directory, job, move |progress, finished| { + heartbeat(&url, progress, finished) + }) } fn start_with_report( directory: &Path, job: Option<&str>, - mut report: impl FnMut(&Progress) + Send + 'static, + mut report: impl FnMut(&Progress, bool) + Send + 'static, ) -> Self { let snapshot = Arc::new(Mutex::new(Progress { phase: if job.is_some() { @@ -120,6 +123,7 @@ impl Monitor { let db = directory.join("history.db"); let job = job.map(str::to_owned); let (stop, receive) = mpsc::channel(); + let (completed, completion) = mpsc::channel(); thread::spawn(move || { let started = Instant::now(); let mut finish = None; @@ -140,7 +144,7 @@ impl Monitor { .into(); } println!("{} ({}s)", progress.line(), started.elapsed().as_secs()); - report(&progress); + report(&progress, finish.is_some()); if finish.is_some() { break; } @@ -150,10 +154,12 @@ impl Monitor { Err(mpsc::RecvTimeoutError::Timeout) => {} } } + let _ = completed.send(()); }); Self { snapshot, stop: Some(stop), + completion, } } pub fn observer(&self) -> impl Fn(ai_hist::CaptureProgress) + 'static { @@ -165,8 +171,16 @@ impl Monitor { p.total_files = capture.total_files; } } - // Sending stop is sufficient: the detached worker exits after its bounded - // request/final report. Capture, delivery and shutdown never join network I/O. + pub fn finish_before_exit(mut self, success: bool, timeout: Duration) { + if let Some(stop) = self.stop.take() { + let _ = stop.send(success); + } + // Setup/--once gets a chance to flush final queue counts, with a strict + // deadline. Ordinary background finish/drop still never waits. + let _ = self.completion.recv_timeout(timeout); + } + // Background/capture callers only signal stop; the worker exits after its + // bounded request/final report without holding up the caller. pub fn finish(mut self, success: bool) { if let Some(stop) = self.stop.take() { let _ = stop.send(success); @@ -178,27 +192,34 @@ impl Drop for Monitor { self.stop.take(); } } -fn heartbeat(url: &str, progress: &Progress) { - // Progress must never acquire the refresh lock or rotate credentials. The - // delivery path owns renewal; skip this heartbeat if its cached token expired. - let Ok(Some(auth)) = cloud::load_selected_auth(Some(url)) else { - return; +fn heartbeat(url: &str, progress: &Progress, finished: bool) { + // Periodic capture updates only read cached credentials. A final upload + // update may renew an idle token, without waiting for the refresh lock. + let token = if finished && progress.phase != "capture_paused" { + match cloud::try_progress_access_token(url, Duration::from_secs(2)) { + Ok(Some(token)) => token, + _ => return, + } + } else { + let Ok(Some(auth)) = cloud::load_selected_auth(Some(url)) else { + return; + }; + let valid = auth + .access_token_expires_at + .as_deref() + .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) + .is_some_and(|expiry| expiry.timestamp() > chrono::Utc::now().timestamp() + 5); + if !valid + || !auth.access_token.starts_with("rth_at_") + || !auth + .access_token + .bytes() + .all(|byte| byte.is_ascii_graphic()) + { + return; + } + auth.access_token }; - let valid = auth - .access_token_expires_at - .as_deref() - .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) - .is_some_and(|expiry| expiry.timestamp() > chrono::Utc::now().timestamp() + 5); - if !valid - || !auth.access_token.starts_with("rth_at_") - || !auth - .access_token - .bytes() - .all(|byte| byte.is_ascii_graphic()) - { - return; - } - let token = auth.access_token; let result = ureq::AgentBuilder::new() .redirects(0) .timeout(Duration::from_secs(3)) @@ -219,10 +240,10 @@ mod tests { for success in [None, Some(false), Some(true)] { let directory = tempfile::tempdir().unwrap(); let (entered, started) = mpsc::channel(); - let (release, blocked) = mpsc::channel(); + let (release, blocked) = mpsc::channel::<()>(); let (finished, done) = mpsc::channel(); let mut first = true; - let monitor = Monitor::start_with_report(directory.path(), None, move |progress| { + let monitor = Monitor::start_with_report(directory.path(), None, move |progress, _| { if first { first = false; entered.send(()).unwrap(); @@ -252,6 +273,28 @@ mod tests { } } + #[test] + fn completion_barrier_waits_for_the_final_report_but_has_a_deadline() { + let directory = tempfile::tempdir().unwrap(); + let (reported, received) = mpsc::channel(); + let monitor = + Monitor::start_with_report(directory.path(), Some("job"), move |_, finished| { + if finished { + reported.send(()).unwrap(); + } + }); + monitor.finish_before_exit(true, Duration::from_secs(1)); + received.try_recv().unwrap(); + let (release, blocked) = mpsc::channel::<()>(); + let monitor = Monitor::start_with_report(directory.path(), Some("job"), move |_, _| { + let _ = blocked.recv(); + }); + let start = Instant::now(); + monitor.finish_before_exit(true, Duration::from_millis(50)); + assert!(start.elapsed() < Duration::from_millis(500)); + drop(release); + } + #[test] fn legacy_null_discovery_state_counts_as_captured() { let directory = tempfile::tempdir().unwrap(); From 61d95195269f2a86df3784ed35ba2054363fc9ba Mon Sep 17 00:00:00 2001 From: Will Washburn Date: Sat, 19 Sep 2026 22:28:00 -0700 Subject: [PATCH 4/6] Require confirmed Cloud progress before one-shot setup succeeds --- .../relayhistory/rust/src/probe/collector.rs | 3 +- .../relayhistory/rust/src/probe/progress.rs | 34 ++++++++++++------- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/plugins/relayhistory/rust/src/probe/collector.rs b/plugins/relayhistory/rust/src/probe/collector.rs index 4050e719..10b898c6 100644 --- a/plugins/relayhistory/rust/src/probe/collector.rs +++ b/plugins/relayhistory/rust/src/probe/collector.rs @@ -166,7 +166,8 @@ pub fn deliver_captured(directory: &Path, config: &Config, before_exit: bool) -> super::progress::Monitor::start(directory, &config.history_url, Some(&config.job_id)); let result = deliver(&db_path, config); if before_exit { - progress.finish_before_exit(result.is_ok(), Duration::from_secs(5)); + let connected = progress.finish_before_exit(result.is_ok(), Duration::from_secs(5)); + ensure!(connected, user_error("Cloud connection could not be confirmed. Check the endpoint and retry; local data remains queued.")); } else { progress.finish(result.is_ok()); } diff --git a/plugins/relayhistory/rust/src/probe/progress.rs b/plugins/relayhistory/rust/src/probe/progress.rs index 0fe97be1..25535b5b 100644 --- a/plugins/relayhistory/rust/src/probe/progress.rs +++ b/plugins/relayhistory/rust/src/probe/progress.rs @@ -95,7 +95,7 @@ impl Progress { pub struct Monitor { snapshot: Arc>, stop: Option>, - completion: mpsc::Receiver<()>, + completion: mpsc::Receiver, } impl Monitor { pub fn start(directory: &Path, history_url: &str, job: Option<&str>) -> Self { @@ -107,7 +107,7 @@ impl Monitor { fn start_with_report( directory: &Path, job: Option<&str>, - mut report: impl FnMut(&Progress, bool) + Send + 'static, + mut report: impl FnMut(&Progress, bool) -> bool + Send + 'static, ) -> Self { let snapshot = Arc::new(Mutex::new(Progress { phase: if job.is_some() { @@ -127,6 +127,7 @@ impl Monitor { thread::spawn(move || { let started = Instant::now(); let mut finish = None; + let mut acknowledged = false; loop { // The next upload monitor now owns progress. Do not send a // delayed final "scanning" heartbeat after successful capture. @@ -144,7 +145,7 @@ impl Monitor { .into(); } println!("{} ({}s)", progress.line(), started.elapsed().as_secs()); - report(&progress, finish.is_some()); + acknowledged = report(&progress, finish.is_some()); if finish.is_some() { break; } @@ -154,7 +155,7 @@ impl Monitor { Err(mpsc::RecvTimeoutError::Timeout) => {} } } - let _ = completed.send(()); + let _ = completed.send(acknowledged); }); Self { snapshot, @@ -171,13 +172,13 @@ impl Monitor { p.total_files = capture.total_files; } } - pub fn finish_before_exit(mut self, success: bool, timeout: Duration) { + pub fn finish_before_exit(mut self, success: bool, timeout: Duration) -> bool { if let Some(stop) = self.stop.take() { let _ = stop.send(success); } // Setup/--once gets a chance to flush final queue counts, with a strict // deadline. Ordinary background finish/drop still never waits. - let _ = self.completion.recv_timeout(timeout); + self.completion.recv_timeout(timeout).unwrap_or(false) } // Background/capture callers only signal stop; the worker exits after its // bounded request/final report without holding up the caller. @@ -192,17 +193,17 @@ impl Drop for Monitor { self.stop.take(); } } -fn heartbeat(url: &str, progress: &Progress, finished: bool) { +fn heartbeat(url: &str, progress: &Progress, finished: bool) -> bool { // Periodic capture updates only read cached credentials. A final upload // update may renew an idle token, without waiting for the refresh lock. let token = if finished && progress.phase != "capture_paused" { match cloud::try_progress_access_token(url, Duration::from_secs(2)) { Ok(Some(token)) => token, - _ => return, + _ => return false, } } else { let Ok(Some(auth)) = cloud::load_selected_auth(Some(url)) else { - return; + return false; }; let valid = auth .access_token_expires_at @@ -216,7 +217,7 @@ fn heartbeat(url: &str, progress: &Progress, finished: bool) { .bytes() .all(|byte| byte.is_ascii_graphic()) { - return; + return false; } auth.access_token }; @@ -227,9 +228,11 @@ fn heartbeat(url: &str, progress: &Progress, finished: bool) { .post(&format!("{url}/v1/onboarding/heartbeat")) .set("Authorization", &format!("Bearer {token}")) .send_json(serde_json::json!({"progress":progress})); - if result.is_err() { + let acknowledged = result.is_ok_and(|response| (200..300).contains(&response.status())); + if !acknowledged { eprintln!("Cloud progress update unavailable; collection continues locally."); } + acknowledged } #[cfg(test)] @@ -251,6 +254,7 @@ mod tests { } else { finished.send(progress.phase.clone()).unwrap(); } + true }); started.recv_timeout(Duration::from_secs(2)).unwrap(); let start = Instant::now(); @@ -282,17 +286,21 @@ mod tests { if finished { reported.send(()).unwrap(); } + true }); - monitor.finish_before_exit(true, Duration::from_secs(1)); + assert!(monitor.finish_before_exit(true, Duration::from_secs(1))); received.try_recv().unwrap(); let (release, blocked) = mpsc::channel::<()>(); let monitor = Monitor::start_with_report(directory.path(), Some("job"), move |_, _| { let _ = blocked.recv(); + true }); let start = Instant::now(); - monitor.finish_before_exit(true, Duration::from_millis(50)); + assert!(!monitor.finish_before_exit(true, Duration::from_millis(50))); assert!(start.elapsed() < Duration::from_millis(500)); drop(release); + let monitor = Monitor::start_with_report(directory.path(), Some("job"), |_, _| false); + assert!(!monitor.finish_before_exit(true, Duration::from_secs(1))); } #[test] From 2dc2b3a11b09512a8e38daf2de460099542b11e8 Mon Sep 17 00:00:00 2001 From: Will Washburn Date: Sat, 19 Sep 2026 22:32:27 -0700 Subject: [PATCH 5/6] Preserve actionable delivery errors when heartbeat confirmation fails --- plugins/relayhistory/rust/src/probe/collector.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/plugins/relayhistory/rust/src/probe/collector.rs b/plugins/relayhistory/rust/src/probe/collector.rs index 10b898c6..618c84e4 100644 --- a/plugins/relayhistory/rust/src/probe/collector.rs +++ b/plugins/relayhistory/rust/src/probe/collector.rs @@ -167,7 +167,9 @@ pub fn deliver_captured(directory: &Path, config: &Config, before_exit: bool) -> let result = deliver(&db_path, config); if before_exit { let connected = progress.finish_before_exit(result.is_ok(), Duration::from_secs(5)); - ensure!(connected, user_error("Cloud connection could not be confirmed. Check the endpoint and retry; local data remains queued.")); + if result.is_ok() { + ensure!(connected, user_error("Cloud connection could not be confirmed. Check the endpoint and retry; local data remains queued.")); + } } else { progress.finish(result.is_ok()); } From dafe88de2933546ffd0927664783ee02d66c04ca Mon Sep 17 00:00:00 2001 From: Will Washburn Date: Sat, 19 Sep 2026 22:42:38 -0700 Subject: [PATCH 6/6] fix(probe): cover bounded heartbeat work in setup wait --- plugins/relayhistory/rust/src/probe/collector.rs | 3 ++- plugins/relayhistory/rust/src/probe/progress.rs | 10 ++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/plugins/relayhistory/rust/src/probe/collector.rs b/plugins/relayhistory/rust/src/probe/collector.rs index 618c84e4..ee79672d 100644 --- a/plugins/relayhistory/rust/src/probe/collector.rs +++ b/plugins/relayhistory/rust/src/probe/collector.rs @@ -166,7 +166,8 @@ pub fn deliver_captured(directory: &Path, config: &Config, before_exit: bool) -> super::progress::Monitor::start(directory, &config.history_url, Some(&config.job_id)); let result = deliver(&db_path, config); if before_exit { - let connected = progress.finish_before_exit(result.is_ok(), Duration::from_secs(5)); + let connected = + progress.finish_before_exit(result.is_ok(), super::progress::COMPLETION_TIMEOUT); if result.is_ok() { ensure!(connected, user_error("Cloud connection could not be confirmed. Check the endpoint and retry; local data remains queued.")); } diff --git a/plugins/relayhistory/rust/src/probe/progress.rs b/plugins/relayhistory/rust/src/probe/progress.rs index 25535b5b..1d4d9bb8 100644 --- a/plugins/relayhistory/rust/src/probe/progress.rs +++ b/plugins/relayhistory/rust/src/probe/progress.rs @@ -9,6 +9,12 @@ use std::{ time::{Duration, Instant}, }; +const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(3); +const REFRESH_TIMEOUT: Duration = Duration::from_secs(2); +// In-flight heartbeat + final token renewal/report + local bookkeeping margin. +pub const COMPLETION_TIMEOUT: Duration = + Duration::from_secs(HEARTBEAT_TIMEOUT.as_secs() * 2 + REFRESH_TIMEOUT.as_secs() + 1); + #[derive(Clone, Default, Serialize)] #[serde(rename_all = "camelCase")] pub struct Progress { @@ -197,7 +203,7 @@ fn heartbeat(url: &str, progress: &Progress, finished: bool) -> bool { // Periodic capture updates only read cached credentials. A final upload // update may renew an idle token, without waiting for the refresh lock. let token = if finished && progress.phase != "capture_paused" { - match cloud::try_progress_access_token(url, Duration::from_secs(2)) { + match cloud::try_progress_access_token(url, REFRESH_TIMEOUT) { Ok(Some(token)) => token, _ => return false, } @@ -223,7 +229,7 @@ fn heartbeat(url: &str, progress: &Progress, finished: bool) -> bool { }; let result = ureq::AgentBuilder::new() .redirects(0) - .timeout(Duration::from_secs(3)) + .timeout(HEARTBEAT_TIMEOUT) .build() .post(&format!("{url}/v1/onboarding/heartbeat")) .set("Authorization", &format!("Bearer {token}"))