Basic Identity system for Tiles

tilekit/Cargo.toml

@@ -12,3 +12,6 @@ tokio = { version = "1" , features = ["macros", "rt-multi-thread"]}
 dspy-rs = "0.7.3"
 bon = "3.7.0"
 indexmap = "2.2.0"
+ucan = { git = "https://github.com/ucan-wg/rs-ucan.git", branch = "main", package = "ucan" }
+ed25519-dalek = { version = "2.2", features = ["rand_core"] }
+keyring = { version = "3", features = ["apple-native"] }

tilekit/src/accounts.rs

@@ -0,0 +1,38 @@
+// Handles stuff related to accounts, identity etc..
+use anyhow::Result;
+use ed25519_dalek::{SigningKey, ed25519::signature::rand_core::OsRng};
+use keyring::Entry;
+use ucan::did::Ed25519Did;
+
+type DID = String;
+type Identity = DID;
+
+/// Creates an `Identity`
+/// # Parameters
+/// - `app`: The service for which Identity is made (for ex: tiles)
+/// The keypair generated will be stored in OS secure storage
+/// Returns an `Identity`
+pub fn create_identity(app: &str) -> Result<Identity> {
+    let mut csprng = OsRng;
+    let signing_key = SigningKey::generate(&mut csprng);
+    let ed_did = Ed25519Did::from(signing_key.clone());
+    let did = ed_did.to_string();
+    let entry = Entry::new(app, &did)?;
+    entry.set_secret(&signing_key.to_keypair_bytes())?;
+    Ok(did)
+}
+
+#[cfg(test)]
+mod tests {
+    use keyring::{mock, set_default_credential_builder};
+
+    use super::*;
+
+    #[test]
+    fn test_create_success() -> Result<()> {
+        set_default_credential_builder(mock::default_credential_builder());
+        let did = create_identity("tiles")?;
+        assert!(did.starts_with("did:key"));
+        Ok(())
+    }
+}

tilekit/src/lib.rs

@@ -1,2 +1,3 @@
+pub mod accounts;
 pub mod modelfile;
 pub mod optimize;

tilekit/src/modelfile.rs

@@ -341,11 +341,9 @@ fn parse_singleline(input: &str) -> IResult<&str, &str> {
     .parse(input)
 }
 fn create_modelfile(commands: Vec<(&str, Output)>) -> Result<Modelfile, String> {
-    // TODO: There might be a better way
     let mut modelfile: Modelfile = Modelfile::new();
     for command in commands {
         let _ = match (command.0.to_lowercase().as_str(), command.1) {
-            //TODO: Can add validations for path if its a gguf file later
             ("from", Output::Single(from)) => modelfile.add_from(from.trim()),
             ("parameter", Output::Pair((param, argument))) => {
                 modelfile.add_parameter(param, argument.trim())

tiles/Cargo.toml

@@ -15,6 +15,7 @@ owo-colors = "4"
 futures-util = "0.3"
 hf-hub = {version = "0.4", features = ["tokio"]}
 rustyline = "17.0"
-
+toml = "1.0.3"
 [dev-dependencies]
 tempfile = "3"
+keyring = { version = "3", features = ["apple-native"] }

tiles/src/commands/mod.rs

@@ -1,12 +1,18 @@
 // Module that handles CLI commands
 
+use anyhow::Result;
 use owo_colors::OwoColorize;
 use tiles::runtime::Runtime;
-use tiles::utils::config::set_memory_path;
+use tiles::utils::accounts::{
+    create_root_account, get_root_user_details, save_root_account, set_nickname,
+};
+use tiles::utils::config::{get_or_create_config, set_memory_path};
 use tiles::{core::health, runtime::RunArgs};
 
 pub use tilekit::optimize::optimize;
 
+use crate::{AccountArgs, AccountCommands};
+
 pub async fn run(runtime: &Runtime, run_args: RunArgs) {
     let _ = runtime.run(run_args).await;
 }
@@ -33,3 +39,54 @@ pub async fn start_server(runtime: &Runtime) {
 pub async fn stop_server(runtime: &Runtime) {
     let _ = runtime.stop_server_daemon().await;
 }
+
+//TODO: add docs
+pub fn run_account_commands(account_args: AccountArgs) -> Result<()> {
+    let config = get_or_create_config()?;
+    let root_user_details = get_root_user_details(&config)?;
+    match account_args.command {
+        Some(AccountCommands::Create { nickname }) => {
+            //TODO: could show a message if did is already there
+            let root_user_config = create_root_account(&config, nickname)?;
+            let id = root_user_config.get("id").unwrap().as_str().unwrap();
+            save_root_account(config, &root_user_config)?;
+            //TODO: color it...
+            let _ = println!("Root account has been created with id: {}", id).green();
+        }
+        Some(AccountCommands::SetNickname { nickname }) => {
+            //FIXME: redundant code
+            if root_user_details.get("id").unwrap().is_empty() {
+                println!(
+                    "Root account not created yet, use {}",
+                    format!("tiles account create").yellow()
+                );
+            } else {
+                match set_nickname(&config, nickname) {
+                    Ok(root_user_config) => {
+                        let id = root_user_config.get("id").unwrap().as_str().unwrap();
+                        let nickname = root_user_config.get("nickname").unwrap().as_str().unwrap();
+                        save_root_account(config, &root_user_config)?;
+                        println!("Nickname {} has been set for ID: {}", nickname, id)
+                    }
+                    Err(err) => {
+                        println!("Failed to set nickname due to {}", err)
+                    }
+                }
+            }
+        }
+        _ => {
+            if root_user_details.get("id").unwrap().is_empty() {
+                println!(
+                    "Root account not created yet, use {}",
+                    format!("tiles account create").yellow()
+                );
+            } else {
+                for elem in root_user_details {
+                    println!("{}: {}", elem.0, elem.1)
+                }
+            }
+        }
+    }
+
+    Ok(())
+}

tiles/src/main.rs

@@ -44,6 +44,8 @@ enum Commands {
         #[arg(long, default_value = "openai:gpt-4o-mini")]
         model: String,
     },
+    /// Manage user account
+    Account(AccountArgs),
 }
 
 #[derive(Debug, Args)]
@@ -90,6 +92,23 @@ enum MemoryCommands {
     SetPath { path: String },
 }
 
+#[derive(Debug, Args)]
+#[command(args_conflicts_with_subcommands = true)]
+#[command(flatten_help = true)]
+struct AccountArgs {
+    #[command(subcommand)]
+    command: Option<AccountCommands>,
+}
+
+#[derive(Debug, Subcommand)]
+enum AccountCommands {
+    /// Creates a local root account
+    Create { nickname: Option<String> },
+
+    /// Sets nickname to local root account
+    SetNickname { nickname: String },
+}
+
 #[tokio::main(flavor = "current_thread")]
 pub async fn main() -> Result<(), Box<dyn Error>> {
     let cli = Cli::parse();
@@ -126,6 +145,9 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
             std::fs::write(&modelfile_path, modelfile.to_string())?;
             println!("Successfully updated {}", modelfile_path);
         }
+        Commands::Account(account_args) => {
+            commands::run_account_commands(account_args)?;
+        }
     }
     Ok(())
 }

tiles/src/runtime/mlx.rs

@@ -118,7 +118,6 @@ impl MLXRuntime {
             .spawn()
             .expect("failed to start server");
 
-        fs::create_dir_all(&config_dir).context("Failed to create config directory")?;
         std::fs::write(pid_file, child.id().to_string()).unwrap();
         println!("Server started with PID {}", child.id());
         Ok(())

tiles/src/utils/accounts.rs

@@ -0,0 +1,197 @@
+// Stuff related to account and identity system
+
+use std::collections::HashMap;
+
+use anyhow::Result;
+use tilekit::accounts::create_identity;
+use toml::Table;
+
+use crate::utils::config::save_config;
+const ROOT_USER_CONFIG_KEY: &str = "root-user";
+
+//TODO: add docs
+pub fn get_root_user_details(config: &Table) -> Result<HashMap<String, String>> {
+    Ok(get_root_account(&config))
+}
+
+fn get_root_account(config: &Table) -> HashMap<String, String> {
+    let root_user = config.get(ROOT_USER_CONFIG_KEY).unwrap();
+    let root_user_table = root_user.as_table().unwrap();
+    let mut root_user_map = HashMap::new();
+    for ele in root_user_table {
+        root_user_map.insert(ele.0.to_string(), ele.1.as_str().unwrap().to_owned());
+    }
+    root_user_map
+}
+
+// Lets return main config table only, let the caller do whatever it wants...
+// TODO: Needed docs
+pub fn create_root_account(config: &Table, nickname: Option<String>) -> Result<Table> {
+    let root_user = config.get(ROOT_USER_CONFIG_KEY).unwrap();
+    let root_user_table = root_user.as_table().unwrap();
+    let did = root_user_table.get("id").unwrap().as_str().unwrap();
+    if did.is_empty() {
+        let root_user_config = create_root_user(root_user_table, nickname)?;
+        Ok(root_user_config)
+    } else {
+        Ok(root_user_table.clone())
+    }
+}
+
+//TODO: docs
+pub fn save_root_account(mut config: Table, root_user_config: &Table) -> Result<()> {
+    config.insert(
+        String::from(ROOT_USER_CONFIG_KEY),
+        toml::Value::Table(root_user_config.clone()),
+    );
+    save_config(&config)
+}
+
+// TODO: add docs
+pub fn set_nickname(config: &Table, nickname: String) -> Result<Table> {
+    let root_user = config.get(ROOT_USER_CONFIG_KEY).unwrap();
+    let mut root_user_table = root_user.as_table().unwrap().clone();
+    let did = root_user_table.get("id").unwrap().as_str().unwrap();
+    if did.is_empty() {
+        Err(anyhow::anyhow!("No Root user available"))
+    } else {
+        root_user_table.insert("id".to_owned(), toml::Value::String(did.to_owned()));
+        root_user_table.insert("nickname".to_owned(), toml::Value::String(nickname));
+        Ok(root_user_table)
+    }
+}
+
+// TODO: add docs
+fn create_root_user(root_user_config: &Table, nickname: Option<String>) -> Result<Table> {
+    // get root user details
+    let mut root_user_table = root_user_config.clone();
+    match create_identity("tiles") {
+        Ok(did) => {
+            root_user_table.insert("id".to_owned(), toml::Value::String(did));
+            if nickname.is_some() {
+                root_user_table.insert(
+                    "nickname".to_owned(),
+                    toml::Value::String(nickname.unwrap()),
+                );
+            }
+            Ok(root_user_table)
+        }
+        Err(err) => Err(err),
+    }
+}
+
+#[cfg(test)]
+
+mod tests {
+    use keyring::{mock, set_default_credential_builder};
+    use toml::Table;
+
+    use crate::utils::accounts::{create_root_account, get_root_account};
+
+    #[test]
+    fn test_get_root_user_details_empty_id() {
+        let config: Table = toml::from_str(
+            r#"
+                [root-user]
+                id = ''
+                nickname = ''
+            "#,
+        )
+        .unwrap();
+        let acc_details = get_root_account(&config);
+        assert!(acc_details.get("id").unwrap().is_empty());
+    }
+
+    #[test]
+    fn test_get_root_user_details_valid_id() {
+        let config: Table = toml::from_str(
+            r#"
+                [root-user]
+                id = 'did:key:xyz'
+                nickname = ''
+            "#,
+        )
+        .unwrap();
+        let acc_details = get_root_account(&config);
+        assert!(acc_details.get("id").unwrap().contains("did:key"));
+    }
+
+    #[test]
+    fn test_create_root_account_but_exists() {
+        let config: Table = toml::from_str(
+            r#"
+                [root-user]
+                id = 'did:key:xyz'
+                nickname = ''
+            "#,
+        )
+        .unwrap();
+        let root_user = create_root_account(&config, None).unwrap();
+
+        assert_eq!(
+            root_user.get("id").unwrap().as_str().unwrap(),
+            "did:key:xyz"
+        );
+    }
+
+    #[test]
+    fn test_create_root_account_new() {
+        set_default_credential_builder(mock::default_credential_builder());
+        let config: Table = toml::from_str(
+            r#"
+                [root-user]
+                id = ''
+                nickname = ''
+            "#,
+        )
+        .unwrap();
+        let root_user = create_root_account(&config, None).unwrap();
+
+        assert_ne!(
+            root_user.get("id").unwrap().as_str().unwrap(),
+            "did:key:xyz"
+        );
+
+        assert!(
+            root_user
+                .get("id")
+                .unwrap()
+                .as_str()
+                .unwrap()
+                .starts_with("did:key")
+        );
+    }
+
+    #[test]
+    fn test_create_root_account_new_w_nickname() {
+        set_default_credential_builder(mock::default_credential_builder());
+        let config: Table = toml::from_str(
+            r#"
+                [root-user]
+                id = ''
+                nickname = ''
+            "#,
+        )
+        .unwrap();
+        let root_user = create_root_account(&config, Some(String::from("madclaws"))).unwrap();
+
+        assert_ne!(
+            root_user.get("id").unwrap().as_str().unwrap(),
+            "did:key:xyz"
+        );
+
+        assert!(
+            root_user
+                .get("id")
+                .unwrap()
+                .as_str()
+                .unwrap()
+                .starts_with("did:key")
+        );
+
+        assert_eq!(
+            root_user.get("nickname").unwrap().as_str().unwrap(),
+            "madclaws"
+        );
+    }
+}