Fix timeline cancellation and add Gemini 3.1 Flash-Lite

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",

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<std::process::Child, String> {
     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());

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<TimelineJobState, String> {
+    let is_running = state
+        .running_timeline_jobs
+        .lock()
+        .map_err(|_| "Failed to lock running timeline jobs".to_string())?
+        .contains(&current.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(&current.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, &current, &job_id).map_err(|e| e.to_string())?;
+    Ok(current)
+}
+
 #[tauri::command]
 pub fn get_chats(state: tauri::State<'_, AppState>) -> Result<Vec<ChatResponse>, 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<TimelineJobState, String> {
     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<String, bool> {
     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<i32
 #[cfg(test)]
 mod smoke_tests {
     use super::*;
+    use std::collections::{HashMap, HashSet};
     use std::path::PathBuf;
+    use std::sync::{Arc, Mutex};
     use std::thread;
     use std::time::{Duration, Instant};
 
@@ -788,6 +824,65 @@ mod smoke_tests {
         }
     }
 
+    fn make_test_state() -> 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() {

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<String> {
+    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<String, String> {
+    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<PathBuf> {
+    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<PathBuf> {
+    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<PathBuf>) -> Vec<PathBuf> {
+    let mut out = Vec::new();
+    for path in paths {
+        if !out.iter().any(|existing| existing == &path) {
+            out.push(path);
+        }
+    }
+    out
+}

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::Wry>) -> tauri::Builder<t
     }
 }
 
-fn load_env_files() {
-    let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
-    let candidates = [cwd.join(".env"), cwd.join("src-tauri").join(".env")];
-
-    for path in candidates {
-        if path.exists() {
-            let _ = dotenvy::from_path(path);
-        }
-    }
-}
-
 fn serve_attachment(
     state: &state::AppState,
     request: &tauri::http::Request<Vec<u8>>,