diff --git a/crates/ai-hist/src/ingest.rs b/crates/ai-hist/src/ingest.rs index 103fd38f..58908f70 100644 --- a/crates/ai-hist/src/ingest.rs +++ b/crates/ai-hist/src/ingest.rs @@ -84,6 +84,69 @@ 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, +} +type CaptureObserver = std::rc::Rc; +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(std::rc::Rc::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| { + let observer = slot.borrow().clone(); + if let Some(observer) = observer { + 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 +482,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 +497,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 +507,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 +520,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 +528,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 +550,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 +1642,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 +2722,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 +4043,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 +4301,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 +4418,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 +4453,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 +8186,93 @@ 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(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); + 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 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(); + 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/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 8ab0eb83..ee79672d 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, false) +} + +pub fn deliver_captured(directory: &Path, config: &Config, before_exit: bool) -> 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,19 @@ 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); + if before_exit { + 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.")); + } + } else { + progress.finish(result.is_ok()); + } + let status = result?; println!( "Probe connected: {} records received, {} queued.", status.acknowledged_records, status.pending_records @@ -250,8 +265,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, false) + }; + if result.is_err() { eprintln!("Sync paused or offline. Retrying; local data remains queued."); } for _ in 0..20 { @@ -380,8 +402,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..e8e26348 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, 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 new file mode 100644 index 00000000..1d4d9bb8 --- /dev/null +++ b/plugins/relayhistory/rust/src/probe/progress.rs @@ -0,0 +1,343 @@ +//! 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}, +}; + +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 { + 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 COALESCE(discovery_state, 'full') = '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 + ), + }, + "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 + ), + "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>, + 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, finished| { + heartbeat(&url, progress, finished) + }) + } + fn start_with_report( + directory: &Path, + job: Option<&str>, + mut report: impl FnMut(&Progress, bool) -> bool + Send + 'static, + ) -> 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 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; + let mut acknowledged = false; + 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 = if job.is_none() { + "capture_paused" + } else { + "paused" + } + .into(); + } + println!("{} ({}s)", progress.line(), started.elapsed().as_secs()); + acknowledged = report(&progress, finish.is_some()); + 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) => {} + } + } + let _ = completed.send(acknowledged); + }); + Self { + snapshot, + stop: Some(stop), + completion, + } + } + 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_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. + 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. + pub fn finish(mut self, success: bool) { + if let Some(stop) = self.stop.take() { + let _ = stop.send(success); + } + } +} +impl Drop for Monitor { + fn drop(&mut self) { + self.stop.take(); + } +} +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, REFRESH_TIMEOUT) { + Ok(Some(token)) => token, + _ => return false, + } + } else { + let Ok(Some(auth)) = cloud::load_selected_auth(Some(url)) else { + return false; + }; + 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 false; + } + auth.access_token + }; + let result = ureq::AgentBuilder::new() + .redirects(0) + .timeout(HEARTBEAT_TIMEOUT) + .build() + .post(&format!("{url}/v1/onboarding/heartbeat")) + .set("Authorization", &format!("Bearer {token}")) + .send_json(serde_json::json!({"progress":progress})); + 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)] +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(); + } + true + }); + 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 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(); + } + true + }); + 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(); + 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] + 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 { + 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")); + } +}