[desktop] Registry support (#459)

src-tauri/Cargo.toml

@@ -15,6 +15,7 @@
   sys-info = "0.9.1"
   sysinfo = "0.29.10"
   thiserror = "1.0.49"
+  tauri-plugin-store = { git = "https://github.com/tauri-apps/plugins-workspace", branch = "v1" }
 
   [dependencies.futures]
     default-features = false

src-tauri/src/controller_binaries.rs

@@ -7,7 +7,7 @@ use crate::{
     errors::{Context, Result},
     logerr,
     swarm::{create_environment, Config},
-    Service, SharedState,
+    utils, Registry, Service, SharedState,
 };
 
 use std::path::PathBuf;
@@ -20,6 +20,7 @@ use sys_info::mem_info;
 use sysinfo::{Pid, ProcessExt, ProcessRefreshKind, RefreshKind, SystemExt};
 
 use tauri::{AppHandle, Runtime, State, Window};
+use tauri_plugin_store::StoreBuilder;
 
 use tokio::process::{Child, Command};
 use tokio::time::interval;
@@ -495,3 +496,143 @@ pub async fn add_service(service: Service, state: State<'_, Arc<SharedState>>) -
     services_guard.insert(service.get_id()?, service);
     Ok(())
 }
+
+#[tauri::command(async)]
+pub async fn add_registry(
+    registry: Registry,
+    app_handle: AppHandle,
+    state: State<'_, Arc<SharedState>>,
+) -> Result<()> {
+    let store_path = app_handle
+        .path_resolver()
+        .app_data_dir()
+        .with_context(|| "Failed to resolve app data dir")?
+        .join("store.json");
+    let mut store = StoreBuilder::new(app_handle, store_path).build();
+    store.load().with_context(|| "Failed to load store")?;
+    if let Some(registries) = store.get("registries").cloned() {
+        match serde_json::from_value::<Vec<Registry>>(registries) {
+            Ok(mut registries) => {
+                registries.push(registry);
+                store
+                    .insert(
+                        "registries".to_string(),
+                        serde_json::to_value(&registries).unwrap(),
+                    )
+                    .with_context(|| "Failed to insert into store")?;
+                utils::fetch_all_services_manifests(&registries, &state)
+                    .await
+                    .expect("failed to fetch services")
+            }
+            Err(e) => println!("Error unwrapping registries: {:?}", e),
+        }
+    } else {
+        let new_registry = [registry];
+        store
+            .insert(
+                "registries".to_string(),
+                serde_json::to_value(&new_registry).unwrap(),
+            )
+            .with_context(|| "Failed to insert into store")?;
+        utils::fetch_all_services_manifests(&new_registry, &state)
+            .await
+            .expect("failed to fetch services")
+    }
+    store.save().expect("failed to save store");
+    Ok(())
+}
+
+#[tauri::command(async)]
+pub async fn delete_registry(
+    registry: Registry,
+    app_handle: AppHandle,
+    state: State<'_, Arc<SharedState>>,
+) -> Result<()> {
+    let store_path = app_handle
+        .path_resolver()
+        .app_data_dir()
+        .with_context(|| "Failed to resolve app data dir")?
+        .join("store.json");
+    let mut store = StoreBuilder::new(app_handle, store_path).build();
+    store.load().with_context(|| "Failed to load store")?;
+    if let Some(registries) = store.get("registries").cloned() {
+        let mut registries = serde_json::from_value::<Vec<Registry>>(registries)
+            .with_context(|| "Failed to deserialize")?;
+        registries.retain(|r| r.url != registry.url);
+        store
+            .insert(
+                "registries".to_string(),
+                serde_json::to_value(registries).with_context(|| "Failed to serialize")?,
+            )
+            .with_context(|| "Failed to insert into store")?;
+        store.save().with_context(|| "Failed to save store")?;
+
+        // Reset services state and refetch all registries
+        let mut services_guard = state.services.lock().await;
+        services_guard.clear();
+        drop(services_guard);
+        if let Some(registries) = store.get("registries").cloned() {
+            match serde_json::from_value::<Vec<Registry>>(registries) {
+                Ok(registries) => utils::fetch_all_services_manifests(&registries, &state)
+                    .await
+                    .with_context(|| "Failed to fetch services")?,
+                Err(e) => log::error!("Error unwrapping registries: {:?}", e),
+            }
+        } else {
+            println!("No registries found");
+        }
+    }
+    Ok(())
+}
+
+#[tauri::command(async)]
+pub async fn fetch_registries(app_handle: AppHandle) -> Result<Vec<Registry>> {
+    let store_path = app_handle
+        .path_resolver()
+        .app_data_dir()
+        .with_context(|| "Failed to resolve app data dir")?
+        .join("store.json");
+    let mut store = StoreBuilder::new(app_handle, store_path).build();
+    match store.load() {
+        Ok(_) => {
+            if let Some(registries) = store.get("registries").cloned() {
+                match serde_json::from_value::<Vec<Registry>>(registries) {
+                    Ok(registries) => Ok(registries),
+                    Err(e) => {
+                        log::error!("Error unwrapping registries: {:?}", e);
+                        Ok(Vec::new())
+                    }
+                }
+            } else {
+                log::error!("No registries found");
+                Ok(Vec::new())
+            }
+        }
+        Err(e) => {
+            log::error!("Error loading store: {:?}", e);
+            Ok(Vec::new())
+        }
+    }
+}
+
+#[tauri::command(async)]
+pub async fn reset_default_registry(
+    app_handle: AppHandle,
+    state: State<'_, Arc<SharedState>>,
+) -> Result<()> {
+    let store_path = app_handle
+        .path_resolver()
+        .app_data_dir()
+        .with_context(|| "Failed to resolve app data dir")?
+        .join("store.json");
+    let mut store = StoreBuilder::new(app_handle.clone(), store_path).build();
+    store.load().with_context(|| "Failed to load store")?;
+    store
+        .delete("registries")
+        .with_context(|| "Failed to delete registries")?;
+    store.save().with_context(|| "Failed to save store")?;
+    add_registry(Registry::default(), app_handle.clone(), state)
+        .await
+        .with_context(|| "Failed to add default registry")?;
+    Ok(())
+}

src-tauri/src/main.rs

@@ -17,6 +17,7 @@ use tauri::{
     AboutMetadata, CustomMenuItem, Manager, Menu, MenuItem, RunEvent, Submenu, SystemTray,
     SystemTrayEvent, SystemTrayMenu, SystemTrayMenuItem, WindowEvent,
 };
+use tauri_plugin_store::StoreBuilder;
 use tokio::process::Child;
 use tokio::sync::Mutex;
 
@@ -28,6 +29,32 @@ pub struct SharedState {
     services: Mutex<HashMap<String, Service>>,
 }
 
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct Registry {
+    url: String,
+}
+
+impl Default for Registry {
+    fn default() -> Self {
+        // Determine the URL based on whether the app is in debug or release mode
+        let url = if cfg!(debug_assertions) {
+            // Debug mode URL
+            "https://raw.githubusercontent.com/premAI-io/prem-registry/dev/manifests.json"
+        } else {
+            // Release mode URL
+            "https://raw.githubusercontent.com/premAI-io/prem-registry/v1/manifests.json"
+        };
+        Registry {
+            url: url.to_string(),
+        }
+    }
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct Store {
+    registries: Vec<Registry>,
+}
+
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct Service {
     // Static state from registry manifest
@@ -193,6 +220,7 @@ fn main() {
 
     let app = tauri::Builder::default()
         .plugin(sentry_tauri::plugin())
+        .plugin(tauri_plugin_store::Builder::default().build())
         .manage(state.clone())
         .invoke_handler(tauri::generate_handler![
             controller_binaries::start_service,
@@ -207,6 +235,10 @@ fn main() {
             controller_binaries::get_service_stats,
             controller_binaries::get_gpu_stats,
             controller_binaries::add_service,
+            controller_binaries::add_registry,
+            controller_binaries::delete_registry,
+            controller_binaries::fetch_registries,
+            controller_binaries::reset_default_registry,
             swarm::is_swarm_supported,
             swarm::get_username,
             swarm::get_petals_models,
@@ -259,21 +291,40 @@ fn main() {
         })
         .setup(|app| {
             tauri::async_runtime::block_on(async move {
-                // Determine the URL based on whether the app is in debug or release mode
-                let url = if cfg!(debug_assertions) {
-                    // Debug mode URL
-                    "https://raw.githubusercontent.com/premAI-io/prem-registry/dev/manifests.json"
-                } else {
-                    // Release mode URL
-                    "https://raw.githubusercontent.com/premAI-io/prem-registry/v1/manifests.json"
-                };
-
-                utils::fetch_services_manifests(
-                    url,
-                    app.state::<Arc<SharedState>>().deref().clone(),
-                )
-                .await
-                .expect("Failed to fetch and save services manifests");
+                //Create a store with default registry if doesn't exist
+                let store_path = app
+                    .path_resolver()
+                    .app_data_dir()
+                    .expect("failed to resolve app data dir")
+                    .join("store.json");
+                if !store_path.exists() {
+                    let mut registries: Vec<Registry> = Vec::new();
+                    registries.push(Registry::default());
+                    let mut default_store = HashMap::new();
+                    default_store.insert(
+                        "registries".to_string(),
+                        serde_json::to_value(registries).unwrap(),
+                    );
+                    let store = StoreBuilder::new(app.handle(), store_path.clone())
+                        .defaults(default_store)
+                        .build();
+                    store.save().expect("failed to save store");
+                    log::info!("Store created");
+                }
+                // Fetch all registries
+                let mut store = StoreBuilder::new(app.handle(), store_path.clone()).build();
+                store.load().expect("Failed to load store");
+                if let Some(registries) = store.get("registries").cloned() {
+                    match serde_json::from_value::<Vec<Registry>>(registries) {
+                        Ok(registries) => utils::fetch_all_services_manifests(
+                            &registries,
+                            &app.state::<Arc<SharedState>>().clone(),
+                        )
+                        .await
+                        .expect("failed to fetch services"),
+                        Err(e) => println!("Error unwrapping registries: {:?}", e),
+                    }
+                }
             });
             Ok(())
         })

src-tauri/src/utils.rs

@@ -1,10 +1,31 @@
 use crate::errors::{Context, Result};
-use crate::{err, Service, SharedState};
+use crate::{err, Registry, Service, SharedState};
+use futures::future;
 use reqwest::get;
 use std::collections::HashMap;
 use std::sync::Arc;
 
-pub async fn fetch_services_manifests(url: &str, state: Arc<SharedState>) -> Result<()> {
+pub async fn fetch_all_services_manifests(
+    registries: &[Registry],
+    state: &Arc<SharedState>,
+) -> Result<()> {
+    let mut handlers = vec![];
+    for registry in registries {
+        let handler = async move {
+            if let Err(err) = fetch_services_manifests(registry.url.as_str(), state).await {
+                println!(
+                    "Failed to fetch {} and save services manifests: {}",
+                    registry.url, err
+                );
+            }
+        };
+        handlers.push(handler);
+    }
+    future::join_all(handlers).await;
+    Ok(())
+}
+
+async fn fetch_services_manifests(url: &str, state: &Arc<SharedState>) -> Result<()> {
     let response = get(url)
         .await
         .with_context(|| format!("Couldn't fetch the manifest from {url:?}"))?;
@@ -13,13 +34,15 @@ pub async fn fetch_services_manifests(url: &str, state: Arc<SharedState>) -> Res
         .await
         .with_context(|| "Failed to parse response to list of services")?;
     let mut services_guard = state.services.lock().await;
-    // TODO: discuss why do we need global ids, why not use uuids and generate them on each load
-    *services_guard = services
-        .into_iter()
-        // removes services without id
-        .filter_map(|x| Some((x.id.clone()?, x)))
-        // removes duplicate services
-        .collect();
+    for service in services {
+        if !service
+            .get_id_ref()
+            .map(|id| services_guard.contains_key(id))
+            .unwrap_or_default()
+        {
+            services_guard.insert(service.get_id()?, service);
+        }
+    }
     Ok(())
 }
 

src/controller/abstractServiceController.ts

@@ -1,4 +1,5 @@
 import type { Service } from "../modules/service/types";
+import type { Registry } from "../modules/settings/types";
 import type { Interface } from "../shared/helpers/interfaces";
 
 import type { DownloadArgs } from "./serviceController";
@@ -23,6 +24,10 @@ abstract class AbstractServiceController {
   abstract getGPUStats(): Promise<Record<string, string>>;
   abstract getInterfaces(): Promise<Interface[]>;
   abstract addService(service: Service): Promise<void>;
+  abstract addRegistry(registry: Registry): Promise<void>;
+  abstract deleteRegistry(registry: Registry): Promise<void>;
+  abstract fetchRegistries(): Promise<Registry[]>;
+  abstract resetDefaultRegistry(): Promise<void>;
 }
 
 export default AbstractServiceController;