From a77a76cef1d4b26754cce1ae1c41fdc099f385d1 Mon Sep 17 00:00:00 2001 From: Luke Whaley Date: Mon, 9 Mar 2026 03:02:34 -0600 Subject: [PATCH 1/2] Fix timeline cancellation and add Gemini 3.1 Flash-Lite --- assistant-agent/run-turn.mjs | 1 + src-tauri/src/assistant_bridge.rs | 2 + src-tauri/src/commands.rs | 121 ++++++++++++++++++++++++++---- src-tauri/src/env_config.rs | 97 ++++++++++++++++++++++++ src-tauri/src/main.rs | 15 +--- src-tauri/src/timeline_ai.rs | 39 +++++----- src-tauri/src/timeline_indexer.rs | 79 +++++++++++++++++-- src/lib/assistant-models.ts | 7 ++ 8 files changed, 308 insertions(+), 53 deletions(-) create mode 100644 src-tauri/src/env_config.rs diff --git a/assistant-agent/run-turn.mjs b/assistant-agent/run-turn.mjs index 7afd7ec..9a98cd2 100644 --- a/assistant-agent/run-turn.mjs +++ b/assistant-agent/run-turn.mjs @@ -97,6 +97,7 @@ const SUPPORTED_MODELS_BY_PROVIDER = { "claude-haiku-4-5", ]), google: new Set([ + "gemini-3.1-flash-lite-preview", "gemini-2.5-flash-lite", "gemini-2.5-flash", "gemini-2.5-pro", diff --git a/src-tauri/src/assistant_bridge.rs b/src-tauri/src/assistant_bridge.rs index 893f627..0ad6551 100644 --- a/src-tauri/src/assistant_bridge.rs +++ b/src-tauri/src/assistant_bridge.rs @@ -381,10 +381,12 @@ fn emit_stream_event(app_handle: &tauri::AppHandle, stream_id: &str, event: &Ass fn spawn_agent_process() -> Result { let cwd = std::env::current_dir().map_err(|e| e.to_string())?; let script_path = resolve_script_path(&cwd)?; + let env_overrides = crate::env_config::assistant_env_overrides(); let mut cmd = Command::new("node"); cmd.arg(script_path) .current_dir(cwd) + .envs(env_overrides) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index fb5e7a5..714a580 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -17,6 +17,46 @@ pub fn start_window_drag(window: tauri::Window) -> Result<(), String> { window.start_dragging().map_err(|e| e.to_string()) } +fn reconcile_stale_timeline_job_state( + state: &AppState, + conn: &rusqlite::Connection, + mut current: TimelineJobState, +) -> Result { + let is_running = state + .running_timeline_jobs + .lock() + .map_err(|_| "Failed to lock running timeline jobs".to_string())? + .contains(¤t.chat_id); + + if is_running || !matches!(current.status.as_str(), "running" | "canceling") { + return Ok(current); + } + + if let Ok(mut canceled) = state.cancel_timeline_jobs.lock() { + canceled.remove(¤t.chat_id); + } + + if current.status == "canceling" { + current.status = "canceled".to_string(); + current.phase = "finalizing".to_string(); + current.error = Some("Canceled by user".to_string()); + } else { + current.status = "failed".to_string(); + current.phase = "failed".to_string(); + current.error = Some("Timeline indexing stopped unexpectedly; previous worker is no longer running".to_string()); + } + + let now = timeline_db::now_iso(); + current.updated_at = Some(now.clone()); + current.finished_at = Some(now); + let job_id = current + .run_id + .clone() + .unwrap_or_else(|| format!("reconciled-{}", Uuid::new_v4())); + timeline_db::set_job_state(conn, ¤t, &job_id).map_err(|e| e.to_string())?; + Ok(current) +} + #[tauri::command] pub fn get_chats(state: tauri::State<'_, AppState>) -> Result, String> { let db = rusqlite::Connection::open_with_flags( @@ -372,6 +412,7 @@ pub fn cancel_timeline_index_impl( let conn = timeline_db::open_rw(&state.timeline_db_path).map_err(|e| e.to_string())?; let mut current = timeline_db::get_job_state(&conn, chat_id).map_err(|e| e.to_string())?; + current = reconcile_stale_timeline_job_state(state, &conn, current)?; let terminal = matches!( current.status.as_str(), @@ -422,7 +463,8 @@ pub fn get_timeline_index_state_impl( ) -> Result { eprintln!("[timeline-cmd] get_state chat_id={}", chat_id); let conn = timeline_db::open_rw(&state.timeline_db_path).map_err(|e| e.to_string())?; - timeline_db::get_job_state(&conn, chat_id).map_err(|e| e.to_string()) + let current = timeline_db::get_job_state(&conn, chat_id).map_err(|e| e.to_string())?; + reconcile_stale_timeline_job_state(state, &conn, current) } #[tauri::command] @@ -730,27 +772,19 @@ pub fn get_assistant_provider_availability() -> HashMap { let mut out = HashMap::new(); out.insert( "openai".to_string(), - std::env::var("OPENAI_API_KEY") - .map(|v| !v.trim().is_empty()) - .unwrap_or(false), + crate::env_config::get_env_var("OPENAI_API_KEY").is_some(), ); out.insert( "anthropic".to_string(), - std::env::var("ANTHROPIC_API_KEY") - .map(|v| !v.trim().is_empty()) - .unwrap_or(false), + crate::env_config::get_env_var("ANTHROPIC_API_KEY").is_some(), ); out.insert( "google".to_string(), - std::env::var("GOOGLE_GENERATIVE_AI_API_KEY") - .map(|v| !v.trim().is_empty()) - .unwrap_or(false), + crate::env_config::get_env_var("GOOGLE_GENERATIVE_AI_API_KEY").is_some(), ); out.insert( "xai".to_string(), - std::env::var("XAI_API_KEY") - .map(|v| !v.trim().is_empty()) - .unwrap_or(false), + crate::env_config::get_env_var("XAI_API_KEY").is_some(), ); out } @@ -774,7 +808,9 @@ fn query_source_max_rowid(db_path: &std::path::Path, chat_id: i32) -> Result AppState { + let timeline_db_path = std::env::temp_dir().join(format!( + "chatpp-timeline-test-{}.db", + Uuid::new_v4() + )); + crate::timeline_db::init_timeline_db(&timeline_db_path).expect("init test timeline db"); + AppState { + db_path: PathBuf::from("/tmp/chat.db"), + timeline_db_path, + handles: HashMap::new(), + chat_participants: HashMap::new(), + contact_names: HashMap::new(), + contact_photos: HashMap::new(), + running_timeline_jobs: Arc::new(Mutex::new(HashSet::new())), + cancel_timeline_jobs: Arc::new(Mutex::new(HashSet::new())), + } + } + + #[test] + fn get_timeline_state_reconciles_stale_canceling_job() { + let state = make_test_state(); + let mut conn = crate::timeline_db::open_rw(&state.timeline_db_path).expect("open timeline db"); + let mut job = TimelineJobState::idle(42); + job.status = "canceling".to_string(); + job.phase = "canceling".to_string(); + job.progress = 0.25; + job.run_id = Some("run-42".to_string()); + crate::timeline_db::set_job_state(&mut conn, &job, "job-42").expect("persist job"); + + let next = get_timeline_index_state_impl(&state, 42).expect("get reconciled state"); + + assert_eq!(next.status, "canceled"); + assert_eq!(next.phase, "finalizing"); + assert_eq!(next.error.as_deref(), Some("Canceled by user")); + assert!(next.finished_at.is_some()); + } + + #[test] + fn get_timeline_state_reconciles_stale_running_job() { + let state = make_test_state(); + let mut conn = crate::timeline_db::open_rw(&state.timeline_db_path).expect("open timeline db"); + let mut job = TimelineJobState::idle(43); + job.status = "running".to_string(); + job.phase = "image-enrichment".to_string(); + job.progress = 0.10; + job.run_id = Some("run-43".to_string()); + crate::timeline_db::set_job_state(&mut conn, &job, "job-43").expect("persist job"); + + let next = get_timeline_index_state_impl(&state, 43).expect("get reconciled state"); + + assert_eq!(next.status, "failed"); + assert_eq!(next.phase, "failed"); + assert_eq!( + next.error.as_deref(), + Some("Timeline indexing stopped unexpectedly; previous worker is no longer running") + ); + assert!(next.finished_at.is_some()); + } + #[test] #[ignore] fn timeline_endpoint_smoke_loop() { diff --git a/src-tauri/src/env_config.rs b/src-tauri/src/env_config.rs new file mode 100644 index 0000000..978874a --- /dev/null +++ b/src-tauri/src/env_config.rs @@ -0,0 +1,97 @@ +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +pub fn apply_env_files() { + for path in candidate_paths() { + if !path.exists() { + continue; + } + let Ok(iter) = dotenvy::from_path_iter(&path) else { + continue; + }; + for item in iter.flatten() { + if std::env::var_os(&item.0).is_none() { + unsafe { + std::env::set_var(&item.0, &item.1); + } + } + } + } +} + +pub fn get_env_var(name: &str) -> Option { + if let Ok(value) = std::env::var(name) { + if !value.trim().is_empty() { + return Some(value); + } + } + + for path in candidate_paths() { + if !path.exists() { + continue; + } + let Ok(iter) = dotenvy::from_path_iter(&path) else { + continue; + }; + for item in iter.flatten() { + if item.0 == name && !item.1.trim().is_empty() { + return Some(item.1); + } + } + } + + None +} + +pub fn assistant_env_overrides() -> HashMap { + const KEYS: &[&str] = &[ + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "GOOGLE_GENERATIVE_AI_API_KEY", + "XAI_API_KEY", + "OPENAI_MODEL", + "OPENAI_MODEL_TIMELINE_TEXT", + "OPENAI_MODEL_TIMELINE_MEDIA", + "TIMELINE_DB_PATH", + "TIMELINE_AI_MOCK", + ]; + + KEYS.iter() + .filter_map(|key| get_env_var(key).map(|value| ((*key).to_string(), value))) + .collect() +} + +fn candidate_paths() -> Vec { + let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); + let mut paths = vec![cwd.join(".env"), cwd.join("src-tauri").join(".env")]; + + if let Some(repo_root) = repo_root_from(&cwd) { + paths.push(repo_root.join(".env")); + paths.push(repo_root.join("src-tauri").join(".env")); + } + + dedupe_paths(paths) +} + +fn repo_root_from(cwd: &Path) -> Option { + if cwd.join("assistant-agent").exists() && cwd.join("src-tauri").exists() { + return Some(cwd.to_path_buf()); + } + + let parent = cwd.parent()?; + if parent.join("assistant-agent").exists() && parent.join("src-tauri").exists() { + return Some(parent.to_path_buf()); + } + + None +} + +fn dedupe_paths(paths: Vec) -> Vec { + let mut out = Vec::new(); + for path in paths { + if !out.iter().any(|existing| existing == &path) { + out.push(path); + } + } + out +} diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs index f4db2be..30596df 100644 --- a/src-tauri/src/main.rs +++ b/src-tauri/src/main.rs @@ -4,6 +4,7 @@ mod assistant_bridge; mod assistant_tools; mod commands; mod db; +mod env_config; mod state; mod timeline_ai; mod timeline_db; @@ -12,11 +13,10 @@ mod timeline_types; mod types; use imessage_database::util::platform::Platform; -use std::path::PathBuf; use tauri::Manager; fn main() { - load_env_files(); + env_config::apply_env_files(); add_platform_plugins(tauri::Builder::default()) .manage(state::init_app_state()) @@ -66,17 +66,6 @@ fn add_platform_plugins(builder: tauri::Builder) -> tauri::Builder>, diff --git a/src-tauri/src/timeline_ai.rs b/src-tauri/src/timeline_ai.rs index a604a3f..90e306a 100644 --- a/src-tauri/src/timeline_ai.rs +++ b/src-tauri/src/timeline_ai.rs @@ -7,6 +7,8 @@ use std::time::{Duration, Instant}; const DEFAULT_OPENAI_MODEL: &str = "gpt-5-nano"; const DEFAULT_TIMEOUT_SECS: u64 = 45; +const DEFAULT_IMAGE_TIMEOUT_SECS: u64 = 20; +const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 10; const MAX_IMAGE_BYTES: usize = 4 * 1024 * 1024; const DEFAULT_IMAGE_MAX_DIMENSION: i32 = 1600; const DEFAULT_IMAGE_JPEG_QUALITY: i32 = 72; @@ -160,29 +162,21 @@ Rules: "; pub fn is_openai_enabled() -> bool { - std::env::var("OPENAI_API_KEY") - .map(|v| !v.trim().is_empty()) - .unwrap_or(false) + crate::env_config::get_env_var("OPENAI_API_KEY").is_some() } pub fn openai_model_default() -> String { - std::env::var("OPENAI_MODEL") - .ok() - .filter(|v| !v.trim().is_empty()) + crate::env_config::get_env_var("OPENAI_MODEL") .unwrap_or_else(|| DEFAULT_OPENAI_MODEL.to_string()) } pub fn openai_model_text() -> String { - std::env::var("OPENAI_MODEL_TIMELINE_TEXT") - .ok() - .filter(|v| !v.trim().is_empty()) + crate::env_config::get_env_var("OPENAI_MODEL_TIMELINE_TEXT") .unwrap_or_else(openai_model_default) } pub fn openai_model_media() -> String { - std::env::var("OPENAI_MODEL_TIMELINE_MEDIA") - .ok() - .filter(|v| !v.trim().is_empty()) + crate::env_config::get_env_var("OPENAI_MODEL_TIMELINE_MEDIA") .unwrap_or_else(openai_model_default) } @@ -297,13 +291,20 @@ pub fn caption_image_file(path: &str, mime_type: &str) -> Result<(String, String describe_image_for_timeline(path, mime_type) } +fn image_timeout_secs() -> u64 { + crate::env_config::get_env_var("TIMELINE_IMAGE_TIMEOUT_SECS") + .and_then(|v| v.parse::().ok()) + .map(|v| v.clamp(5, 120)) + .unwrap_or(DEFAULT_IMAGE_TIMEOUT_SECS) +} + pub fn describe_image_for_timeline( path: &str, mime_type: &str, ) -> Result<(String, String), String> { let started = Instant::now(); - let api_key = std::env::var("OPENAI_API_KEY") - .map_err(|_| "OPENAI_API_KEY is not set; skipping media captioning".to_string())?; + let api_key = crate::env_config::get_env_var("OPENAI_API_KEY") + .ok_or_else(|| "OPENAI_API_KEY is not set; skipping media captioning".to_string())?; let (bytes, effective_mime, transformed_path) = prepare_image_bytes_for_caption(path, mime_type)?; @@ -315,7 +316,8 @@ pub fn describe_image_for_timeline( }; let client = Client::builder() - .timeout(Duration::from_secs(DEFAULT_TIMEOUT_SECS)) + .connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS)) + .timeout(Duration::from_secs(image_timeout_secs())) .build() .map_err(|e| format!("Failed to initialize HTTP client: {}", e))?; @@ -477,8 +479,7 @@ fn run_responses_call( schema: &Value, ) -> Result { let started = Instant::now(); - if std::env::var("TIMELINE_AI_MOCK") - .ok() + if crate::env_config::get_env_var("TIMELINE_AI_MOCK") .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) .unwrap_or(false) { @@ -493,8 +494,8 @@ fn run_responses_call( return Ok(mock); } - let api_key = std::env::var("OPENAI_API_KEY") - .map_err(|_| "OPENAI_API_KEY is not set; cannot run AI timeline indexing".to_string())?; + let api_key = crate::env_config::get_env_var("OPENAI_API_KEY") + .ok_or_else(|| "OPENAI_API_KEY is not set; cannot run AI timeline indexing".to_string())?; let client = Client::builder() .timeout(Duration::from_secs(DEFAULT_TIMEOUT_SECS)) diff --git a/src-tauri/src/timeline_indexer.rs b/src-tauri/src/timeline_indexer.rs index 242ed3f..2c66a4a 100644 --- a/src-tauri/src/timeline_indexer.rs +++ b/src-tauri/src/timeline_indexer.rs @@ -17,6 +17,7 @@ use rusqlite::Connection; use std::cmp::{max, min}; use std::collections::{HashMap, HashSet, VecDeque}; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::thread; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; @@ -215,7 +216,7 @@ fn run_timeline_index_job_inner( job.started_at = Some(timeline_db::now_iso()); job.updated_at = Some(timeline_db::now_iso()); job.run_id = Some(run_id.clone()); - timeline_db::set_job_state(&timeline_conn, &job, &run_id) + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to initialize timeline job state: {}", e))?; let chat_inputs = load_chat_inputs(&source_conn, config.chat_id, contact_names) @@ -231,7 +232,7 @@ fn run_timeline_index_job_inner( job.total_messages = messages.len() as i32; job.processed_messages = 0; - timeline_db::set_job_state(&timeline_conn, &job, &run_id) + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to set totals: {}", e))?; let image_workers = timeline_image_workers(); @@ -242,10 +243,11 @@ fn run_timeline_index_job_inner( job.phase = "image-enrichment".to_string(); job.progress = 0.10; job.updated_at = Some(timeline_db::now_iso()); - timeline_db::set_job_state(&timeline_conn, &job, &run_id) + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to set image phase: {}", e))?; if timeline_ai::is_openai_enabled() { + let mut progress_persist_error: Option = None; caption_by_attachment = enrich_images_concurrent( &source_conn, source_db_path, @@ -254,7 +256,28 @@ fn run_timeline_index_job_inner( image_retries, cancel_jobs, config.chat_id, + &mut |done, total| { + if total == 0 { + return; + } + job.phase = "image-enrichment".to_string(); + job.processed_messages = done.min(i32::MAX as usize) as i32; + job.total_messages = total.min(i32::MAX as usize) as i32; + job.progress = 0.10 + ((done as f32) / (total as f32)) * 0.08; + job.updated_at = Some(timeline_db::now_iso()); + if let Err(err) = persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) { + if progress_persist_error.is_none() { + progress_persist_error = Some(format!( + "Failed to persist image enrichment progress: {}", + err + )); + } + } + }, ); + if let Some(err) = progress_persist_error { + return Err(err); + } } if is_canceled(cancel_jobs, config.chat_id) { @@ -294,7 +317,7 @@ fn run_timeline_index_job_inner( job.phase = "l0-generation".to_string(); job.progress = 0.18; job.updated_at = Some(timeline_db::now_iso()); - timeline_db::set_job_state(&timeline_conn, &job, &run_id) + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to set l0 phase: {}", e))?; let mut seen_ranges = HashSet::<(i32, i32)>::new(); @@ -463,7 +486,7 @@ fn run_timeline_index_job_inner( job.completed_batches = completed_batches; job.progress = 0.18 + ((window_idx as f32 + 1.0) / (windows.len().max(1) as f32)) * 0.42; job.updated_at = Some(timeline_db::now_iso()); - timeline_db::set_job_state(&timeline_conn, &job, &run_id) + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to persist l0 progress: {}", e))?; } @@ -505,7 +528,7 @@ fn run_timeline_index_job_inner( job.phase = "l2-topics".to_string(); job.progress = 0.66; job.updated_at = Some(timeline_db::now_iso()); - timeline_db::set_job_state(&timeline_conn, &job, &run_id) + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to persist l2 phase: {}", e))?; let l0_nodes_snapshot = all_nodes.clone(); @@ -550,7 +573,7 @@ fn run_timeline_index_job_inner( job.phase = "l1-subtopics".to_string(); job.progress = 0.82; job.updated_at = Some(timeline_db::now_iso()); - timeline_db::set_job_state(&timeline_conn, &job, &run_id) + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to persist l1 phase: {}", e))?; let subtopic_result = build_l1_contiguous_subtopics( @@ -584,7 +607,7 @@ fn run_timeline_index_job_inner( job.phase = "persist".to_string(); job.progress = 0.93; job.updated_at = Some(timeline_db::now_iso()); - timeline_db::set_job_state(&timeline_conn, &job, &run_id) + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to persist persist phase: {}", e))?; let mut evidence = Vec::::new(); @@ -951,6 +974,7 @@ fn enrich_images_concurrent( retries: usize, cancel_jobs: &Arc>>, chat_id: i32, + on_progress: &mut dyn FnMut(usize, usize), ) -> HashMap { let mut tasks = VecDeque::::new(); let mut seen = HashSet::::new(); @@ -976,8 +1000,12 @@ fn enrich_images_concurrent( return HashMap::new(); } + let total_tasks = tasks.len(); + on_progress(0, total_tasks); + let queue = Arc::new(Mutex::new(tasks)); let out = Arc::new(Mutex::new(HashMap::::new())); + let completed = Arc::new(AtomicUsize::new(0)); let worker_count = workers.max(1).min(12); let mut handles = Vec::new(); @@ -986,6 +1014,7 @@ fn enrich_images_concurrent( let queue = queue.clone(); let out = out.clone(); let cancel_jobs = cancel_jobs.clone(); + let completed = completed.clone(); handles.push(thread::spawn(move || loop { if is_canceled(&cancel_jobs, chat_id) { break; @@ -1007,13 +1036,29 @@ fn enrich_images_concurrent( locked.insert(task.attachment_rowid, result); } } + completed.fetch_add(1, Ordering::Relaxed); })); } + let mut last_reported = usize::MAX; + loop { + let done = completed.load(Ordering::Relaxed).min(total_tasks); + if done != last_reported { + on_progress(done, total_tasks); + last_reported = done; + } + if handles.iter().all(|handle| handle.is_finished()) { + break; + } + thread::sleep(Duration::from_millis(500)); + } + for handle in handles { let _ = handle.join(); } + on_progress(completed.load(Ordering::Relaxed).min(total_tasks), total_tasks); + out.lock().map(|v| v.clone()).unwrap_or_default() } @@ -2658,6 +2703,24 @@ fn is_canceled(cancel_jobs: &Arc>>, chat_id: i32) -> bool { .unwrap_or(false) } +fn persist_job_state( + timeline_conn: &Connection, + job: &mut TimelineJobState, + job_id: &str, + cancel_jobs: &Arc>>, +) -> Result<(), String> { + if matches!(job.status.as_str(), "running" | "canceling") + && is_canceled(cancel_jobs, job.chat_id) + { + job.status = "canceling".to_string(); + job.phase = "canceling".to_string(); + job.error = Some("Cancel requested".to_string()); + job.updated_at = Some(timeline_db::now_iso()); + } + + timeline_db::set_job_state(timeline_conn, job, job_id).map_err(|e| e.to_string()) +} + fn mark_canceled( timeline_conn: &Connection, job_id: &str, diff --git a/src/lib/assistant-models.ts b/src/lib/assistant-models.ts index 0b6eb3c..6baf47b 100644 --- a/src/lib/assistant-models.ts +++ b/src/lib/assistant-models.ts @@ -65,6 +65,13 @@ export const ASSISTANT_MODEL_OPTIONS: AssistantModelOption[] = [ providerLabel: "Google", requiredEnvVar: "GOOGLE_GENERATIVE_AI_API_KEY", }, + { + id: "gemini-3.1-flash-lite-preview", + label: "Gemini 3.1 Flash-Lite (Preview)", + provider: "google", + providerLabel: "Google", + requiredEnvVar: "GOOGLE_GENERATIVE_AI_API_KEY", + }, { id: "gemini-2.5-flash", label: "Gemini 2.5 Flash", From 17aecb305d4654dbfce229ec79cadc86a836b443 Mon Sep 17 00:00:00 2001 From: Luke Whaley Date: Mon, 9 Mar 2026 03:10:02 -0600 Subject: [PATCH 2/2] Restore message progress after image enrichment --- src-tauri/src/timeline_indexer.rs | 89 ++++++++++++++++++------------- 1 file changed, 52 insertions(+), 37 deletions(-) diff --git a/src-tauri/src/timeline_indexer.rs b/src-tauri/src/timeline_indexer.rs index 2c66a4a..c1a29c8 100644 --- a/src-tauri/src/timeline_indexer.rs +++ b/src-tauri/src/timeline_indexer.rs @@ -235,48 +235,53 @@ fn run_timeline_index_job_inner( persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) .map_err(|e| format!("Failed to set totals: {}", e))?; - let image_workers = timeline_image_workers(); - let image_retries = timeline_image_retries(); let mut media_insights = Vec::::new(); let mut caption_by_attachment = HashMap::::new(); - job.phase = "image-enrichment".to_string(); - job.progress = 0.10; - job.updated_at = Some(timeline_db::now_iso()); - persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) - .map_err(|e| format!("Failed to set image phase: {}", e))?; + if timeline_image_enrichment_enabled() { + let image_workers = timeline_image_workers(); + let image_retries = timeline_image_retries(); - if timeline_ai::is_openai_enabled() { - let mut progress_persist_error: Option = None; - caption_by_attachment = enrich_images_concurrent( - &source_conn, - source_db_path, - &messages, - image_workers, - image_retries, - cancel_jobs, - config.chat_id, - &mut |done, total| { - if total == 0 { - return; - } - job.phase = "image-enrichment".to_string(); - job.processed_messages = done.min(i32::MAX as usize) as i32; - job.total_messages = total.min(i32::MAX as usize) as i32; - job.progress = 0.10 + ((done as f32) / (total as f32)) * 0.08; - job.updated_at = Some(timeline_db::now_iso()); - if let Err(err) = persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) { - if progress_persist_error.is_none() { - progress_persist_error = Some(format!( - "Failed to persist image enrichment progress: {}", - err - )); + job.phase = "image-enrichment".to_string(); + job.progress = 0.10; + job.updated_at = Some(timeline_db::now_iso()); + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) + .map_err(|e| format!("Failed to set image phase: {}", e))?; + + if timeline_ai::is_openai_enabled() { + let mut progress_persist_error: Option = None; + caption_by_attachment = enrich_images_concurrent( + &source_conn, + source_db_path, + &messages, + image_workers, + image_retries, + cancel_jobs, + config.chat_id, + &mut |done, total| { + if total == 0 { + return; } - } - }, - ); - if let Some(err) = progress_persist_error { - return Err(err); + job.phase = "image-enrichment".to_string(); + job.processed_messages = done.min(i32::MAX as usize) as i32; + job.total_messages = total.min(i32::MAX as usize) as i32; + job.progress = 0.10 + ((done as f32) / (total as f32)) * 0.08; + job.updated_at = Some(timeline_db::now_iso()); + if let Err(err) = + persist_job_state(&timeline_conn, &mut job, &run_id, cancel_jobs) + { + if progress_persist_error.is_none() { + progress_persist_error = Some(format!( + "Failed to persist image enrichment progress: {}", + err + )); + } + } + }, + ); + if let Some(err) = progress_persist_error { + return Err(err); + } } } @@ -314,6 +319,10 @@ fn run_timeline_index_job_inner( let mut openai_used = false; let mut degraded = false; + // Image enrichment temporarily repurposes the progress counters for attachment tasks. + // Restore message-based totals before moving into L0 window generation. + job.processed_messages = 0; + job.total_messages = messages.len().min(i32::MAX as usize) as i32; job.phase = "l0-generation".to_string(); job.progress = 0.18; job.updated_at = Some(timeline_db::now_iso()); @@ -2651,6 +2660,12 @@ fn timeline_image_retries() -> usize { .unwrap_or(DEFAULT_IMAGE_RETRIES) } +fn timeline_image_enrichment_enabled() -> bool { + crate::env_config::get_env_var("TIMELINE_ENABLE_IMAGE_ENRICHMENT") + .map(|v| !matches!(v.to_ascii_lowercase().as_str(), "0" | "false" | "no" | "off")) + .unwrap_or(true) +} + fn timeline_subtopic_max_moments() -> usize { std::env::var("TIMELINE_SUBTOPIC_MAX_MOMENTS") .ok()