Newer
Older
TelosDB / src-tauri / src / lib.rs
mod db;
mod mcp;

use tauri::Manager;
use tauri_plugin_shell::ShellExt;
use tauri_plugin_shell::process::{CommandEvent, CommandChild};
use std::sync::{Arc, Mutex};


#[allow(dead_code)]
struct AppState {
    db_pool: sqlx::SqlitePool,
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    let llama_child: Arc<Mutex<Option<CommandChild>>> = Arc::new(Mutex::new(None));
    tauri::Builder::default()
        .plugin(tauri_plugin_shell::init())
        .plugin({
            let log_dir = if cfg!(debug_assertions) {
                std::env::current_dir().unwrap().join("logs")
            } else {
                dirs::data_dir().unwrap_or_else(|| std::env::current_dir().unwrap()).join("com.telosdb.app").join("logs")
            };
            std::fs::create_dir_all(&log_dir).ok();
            let log_file_path = log_dir.join("telos.log");

            tauri_plugin_log::Builder::default()
                .targets([
                    tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::Stdout),
                    tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::LogDir {
                        file_name: Some("telos.log".to_string()),
                    }),
                    tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::Webview),
                ])
                .rotation_strategy(tauri_plugin_log::RotationStrategy::KeepAll)
                .max_file_size(10 * 1024 * 1024)
                .build()
        })

        .setup({
            let llama_child = llama_child.clone();
            move |app| {
                // Resolve paths
                let app_data_dir = app.path().app_data_dir().expect("failed to get app data dir");
                let db_path = app_data_dir.join("telos.db");

                // llama-serverの起動をTauriのsidecar APIで行う
                let resource_dir = app.path().resource_dir().unwrap_or_default();
                let bin_dir = resource_dir.join("bin");
                let model_path = bin_dir.join("gemma-3-270m-it-Q4_K_M.gguf");
                
                // vec0.dll はビルド時に実行ルート(resource_dir)にコピーされる
                let mut vec0_path = resource_dir.join("vec0.dll");
                
                // 開発時 (target/debug) かつ bin にある場合のフォールバック
                if !vec0_path.exists() && bin_dir.join("vec0.dll").exists() {
                    vec0_path = bin_dir.join("vec0.dll");
                }

                log::info!("Initializing TelosDB at {:?}", db_path);
                log::info!("Bin directory: {:?}", bin_dir);
                log::info!("Model path: {:?}", model_path);
                log::info!("vec0.dll path: {:?}", vec0_path);

                if !vec0_path.exists() {
                    log::error!("vec0.dll NOT FOUND at {:?}. Vector search and DB init will fail.", vec0_path);
                }

                // llama-server自動起動(Tauri sidecar API使用)
                if model_path.exists() {
                    let (mut rx, child) = app.shell()
                        .sidecar("llama-server")
                        .expect("failed to create sidecar")
                        .args(["--model", model_path.to_str().unwrap(), "--port", "8080", "--embedding", "--parallel", "1"])
                        .spawn()
                        .expect("failed to spawn sidecar");

                    log::info!("llama-server sidecar started");
                    *llama_child.lock().unwrap() = Some(child);
                    std::thread::spawn(move || {
                        while let Some(event) = rx.blocking_recv() {
                            match event {
                                CommandEvent::Stdout(line) => log::info!("llama-server: {}", String::from_utf8_lossy(&line)),
                                CommandEvent::Stderr(line) => log::error!("llama-server: {}", String::from_utf8_lossy(&line)),
                                _ => {}
                            }
                        }
                    });
                } else {
                    log::error!("gemma-3-270m-it-Q4_K_M.gguf not found at {:?}", model_path);
                }

                // DB初期化
                match db::init_db(&db_path, &vec0_path) {
                    Ok(_) => log::info!("Database schema initialized."),
                    Err(e) => log::error!("Database schema init failed: {:?}", e),
                }

                // Async Pool
                let pool: sqlx::SqlitePool = tauri::async_runtime::block_on(async {
                    match db::init_pool(db_path.to_str().unwrap(), vec0_path.to_str().unwrap().to_owned()).await {
                        Ok(pool) => {
                            log::info!("App State managed with SQLx pool.");
                            pool
                        },
                        Err(e) => {
                            log::error!("Failed to create SQLx pool: {:?}", e);
                            panic!("Failed to create SQLx pool");
                        }
                    }
                });
                app.manage(AppState { db_pool: pool.clone() });

                // MCP Server
                let db_path_str = db_path.to_str().unwrap().to_owned();
                let vec0_path_str = vec0_path.to_str().unwrap().to_owned();
                
                use tokio::sync::RwLock;
                let llama_status = Arc::new(RwLock::new("unknown".to_string()));
                tauri::async_runtime::spawn({
                    let llama_status = llama_status.clone();
                    async move {
                        mcp::run_server(3001, &db_path_str, &vec0_path_str, llama_status).await;
                    }
                });

                Ok(())
            }
        })
        .on_window_event({
            let llama_child = llama_child.clone();
            move |_app_handle, event| {
                if let tauri::WindowEvent::CloseRequested { .. } = event {
                    // llama-serverプロセスをkill
                    if let Some(child) = llama_child.lock().unwrap().take() {
                        let _ = child.kill();
                    }
                }
            }
        })
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}