feat: add Hub authentication for future sync + extra features (#3010)

5 files changed, 268 insertions, 3 deletions

diff --git a/crates/atuin-client/src/hub.rs b/crates/atuin-client/src/hub.rs
new file mode 100644
index 00000000..5b34574b
--- /dev/null
+++ b/crates/atuin-client/src/hub.rs
@@ -0,0 +1,235 @@
+//! Hub authentication support for Atuin
+//!
+//! This module provides programmatic access to the Atuin Hub authentication flow.
+//! It can be used by other crates (like atuin-ai) to authenticate with the Hub
+//! and obtain session tokens.
+//!
+//! Hub authentication is separate from sync authentication - users can have both
+//! a sync session (for history sync) and a hub session (for Hub-specific features
+//! like AI).
+
+use std::time::Duration;
+
+use eyre::{Context, Result, bail};
+use reqwest::{StatusCode, Url, header::USER_AGENT};
+
+use atuin_common::{
+    api::{
+        ATUIN_CARGO_VERSION, ATUIN_HEADER_VERSION, CliCodeResponse, CliVerifyResponse,
+        ErrorResponse,
+    },
+    tls::ensure_crypto_provider,
+};
+
+use crate::settings::Settings;
+
+static APP_USER_AGENT: &str = concat!("atuin/", env!("CARGO_PKG_VERSION"));
+
+/// The result of starting a hub authentication flow
+#[derive(Debug, Clone)]
+pub struct HubAuthSession {
+    /// The code to be verified
+    pub code: String,
+    /// The URL the user should visit to authenticate
+    pub auth_url: String,
+    /// The hub address being used
+    pub hub_address: String,
+}
+
+/// The result of polling for hub auth completion
+#[derive(Debug, Clone)]
+pub enum HubAuthStatus {
+    /// Still waiting for user authorization
+    Pending,
+    /// Authorization complete, contains the session token
+    Complete(String),
+    /// Authorization failed with an error
+    Failed(String),
+}
+
+/// Default poll interval for checking auth status
+pub const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(2);
+
+/// Default timeout for the entire auth flow
+pub const DEFAULT_AUTH_TIMEOUT: Duration = Duration::from_secs(600);
+
+impl HubAuthSession {
+    /// Start a new hub authentication session
+    ///
+    /// Returns a session containing the code and auth URL that the user should visit.
+    pub async fn start(settings: &Settings) -> Result<Self> {
+        let code_response = request_code(&settings.hub_address)
+            .await
+            .context("Failed to request authentication code from Hub")?;
+
+        let code = code_response.code;
+        let auth_url = format!("{}/auth/cli?code={}", settings.hub_address, code);
+
+        Ok(Self {
+            code,
+            auth_url,
+            hub_address: settings.hub_address.clone(),
+        })
+    }
+
+    /// Poll for the authentication status
+    ///
+    /// Returns the current status of the authentication flow.
+    pub async fn poll(&self) -> Result<HubAuthStatus> {
+        match verify_code(&self.hub_address, &self.code).await {
+            Ok(response) => {
+                if let Some(token) = response.token {
+                    Ok(HubAuthStatus::Complete(token))
+                } else if let Some(error) = response.error {
+                    Ok(HubAuthStatus::Failed(error))
+                } else {
+                    Ok(HubAuthStatus::Pending)
+                }
+            }
+            Err(e) => {
+                // Transient errors shouldn't fail the whole flow
+                log::debug!("Verification poll failed: {}", e);
+                Ok(HubAuthStatus::Pending)
+            }
+        }
+    }
+
+    /// Poll until completion or timeout
+    ///
+    /// This is a convenience method that polls repeatedly until the auth completes
+    /// or times out.
+    pub async fn wait_for_completion(
+        &self,
+        timeout: Duration,
+        poll_interval: Duration,
+    ) -> Result<String> {
+        let start = std::time::Instant::now();
+
+        loop {
+            if start.elapsed() > timeout {
+                bail!("Authentication timed out. Please try again.");
+            }
+
+            match self.poll().await? {
+                HubAuthStatus::Complete(token) => return Ok(token),
+                HubAuthStatus::Failed(error) => bail!("Authentication failed: {}", error),
+                HubAuthStatus::Pending => {
+                    tokio::time::sleep(poll_interval).await;
+                }
+            }
+        }
+    }
+}
+
+/// Save a hub session token
+///
+/// This saves the token to the meta store so it can be used for subsequent Hub API calls.
+/// Note: This is separate from the sync session token.
+pub async fn save_session(token: &str) -> Result<()> {
+    Settings::meta_store()
+        .await?
+        .save_hub_session(token)
+        .await
+        .context("Failed to save hub session")
+}
+
+/// Delete the hub session token (logout from Hub)
+pub async fn delete_session() -> Result<()> {
+    Settings::meta_store()
+        .await?
+        .delete_hub_session()
+        .await
+        .context("Failed to delete hub session")
+}
+
+/// Check if the user is logged in with Hub authentication
+///
+/// Returns true if the user has a valid Hub session token.
+/// This is independent of whether they have a sync session.
+pub async fn is_logged_in() -> Result<bool> {
+    Settings::meta_store().await?.hub_logged_in().await
+}
+
+/// Get the hub session token if available
+///
+/// Returns the Hub session token if the user is logged in with Hub auth,
+/// or None if not logged in.
+pub async fn get_session_token() -> Result<Option<String>> {
+    Settings::meta_store().await?.hub_session_token().await
+}
+
+// --- Internal HTTP functions ---
+
+fn make_url(address: &str, path: &str) -> Result<String> {
+    let address = if address.ends_with('/') {
+        address.to_string()
+    } else {
+        format!("{address}/")
+    };
+
+    let path = path.strip_prefix('/').unwrap_or(path);
+
+    let url = Url::parse(&address)
+        .context("failed to parse hub address")?
+        .join(path)
+        .context("failed to join hub URL path")?;
+
+    Ok(url.to_string())
+}
+
+async fn handle_resp_error(resp: reqwest::Response) -> Result<reqwest::Response> {
+    let status = resp.status();
+
+    if status == StatusCode::SERVICE_UNAVAILABLE {
+        bail!("Service unavailable: check https://status.atuin.sh");
+    }
+
+    if status == StatusCode::TOO_MANY_REQUESTS {
+        bail!("Rate limited; please wait before trying again");
+    }
+
+    if !status.is_success() {
+        if let Ok(error) = resp.json::<ErrorResponse>().await {
+            bail!("Hub error: {} - {}", status, error.reason);
+        }
+        bail!("Hub request failed with status: {}", status);
+    }
+
+    Ok(resp)
+}
+
+/// Request a CLI auth code from the Atuin Hub
+async fn request_code(address: &str) -> Result<CliCodeResponse> {
+    ensure_crypto_provider();
+    let url = make_url(address, "/auth/cli/code")?;
+    let client = reqwest::Client::new();
+
+    let resp = client
+        .post(&url)
+        .header(USER_AGENT, APP_USER_AGENT)
+        .header(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION)
+        .send()
+        .await?;
+    let resp = handle_resp_error(resp).await?;
+
+    let code_response = resp.json::<CliCodeResponse>().await?;
+    Ok(code_response)
+}
+
+/// Poll to verify the CLI auth code and get the session token
+async fn verify_code(address: &str, code: &str) -> Result<CliVerifyResponse> {
+    ensure_crypto_provider();
+    let url = make_url(address, &format!("/auth/cli/verify?code={}", code))?;
+    let client = reqwest::Client::new();
+
+    let resp = client
+        .post(&url)
+        .header(USER_AGENT, APP_USER_AGENT)
+        .header(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION)
+        .send()
+        .await?;
+    let resp = handle_resp_error(resp).await?;
+
+    let verify_response = resp.json::<CliVerifyResponse>().await?;
+    Ok(verify_response)
+}
diff --git a/crates/atuin-client/src/lib.rs b/crates/atuin-client/src/lib.rs
index 160d4529..352e5746 100644
--- a/crates/atuin-client/src/lib.rs
+++ b/crates/atuin-client/src/lib.rs
@@ -5,6 +5,12 @@ extern crate log;
 
 #[cfg(feature = "sync")]
 pub mod api_client;
+#[cfg(feature = "hub")]
+pub mod hub;
+#[cfg(feature = "sync")]
+pub mod login;
+#[cfg(feature = "sync")]
+pub mod register;
 #[cfg(feature = "sync")]
 pub mod sync;
 
@@ -12,13 +18,11 @@ pub mod database;
 pub mod encryption;
 pub mod history;
 pub mod import;
-pub mod login;
 pub mod logout;
 pub mod meta;
 pub mod ordering;
 pub mod plugin;
 pub mod record;
-pub mod register;
 pub mod secrets;
 pub mod settings;
 pub mod theme;
diff --git a/crates/atuin-client/src/meta.rs b/crates/atuin-client/src/meta.rs
index 870f36d0..94ddcaf3 100644
--- a/crates/atuin-client/src/meta.rs
+++ b/crates/atuin-client/src/meta.rs
@@ -21,6 +21,7 @@ const KEY_LAST_SYNC: &str = "last_sync_time";
 const KEY_LAST_VERSION_CHECK: &str = "last_version_check_time";
 const KEY_LATEST_VERSION: &str = "latest_version";
 const KEY_SESSION: &str = "session";
+const KEY_HUB_SESSION: &str = "hub_session";
 const KEY_FILES_MIGRATED: &str = "files_migrated";
 
 pub struct MetaStore {
@@ -189,6 +190,24 @@ impl MetaStore {
         Ok(self.session_token().await?.is_some())
     }
 
+    // Hub session methods (separate from sync session, used for Hub-specific features like AI)
+
+    pub async fn hub_session_token(&self) -> Result<Option<String>> {
+        self.get(KEY_HUB_SESSION).await
+    }
+
+    pub async fn save_hub_session(&self, token: &str) -> Result<()> {
+        self.set(KEY_HUB_SESSION, token).await
+    }
+
+    pub async fn delete_hub_session(&self) -> Result<()> {
+        self.delete(KEY_HUB_SESSION).await
+    }
+
+    pub async fn hub_logged_in(&self) -> Result<bool> {
+        Ok(self.hub_session_token().await?.is_some())
+    }
+
     // File migration: on first open, migrate old plain-text files into the database.
     // Old files are left in place for safe downgrades.
 
diff --git a/crates/atuin-client/src/register.rs b/crates/atuin-client/src/register.rs
index b0c80dc4..ad077dd1 100644
--- a/crates/atuin-client/src/register.rs
+++ b/crates/atuin-client/src/register.rs
@@ -2,7 +2,7 @@ use eyre::Result;
 
 use crate::{api_client, settings::Settings};
 
-pub async fn register(
+pub async fn register_classic(
     settings: &Settings,
     username: String,
     email: String,
diff --git a/crates/atuin-client/src/settings.rs b/crates/atuin-client/src/settings.rs
index 270fc200..7e062e75 100644
--- a/crates/atuin-client/src/settings.rs
+++ b/crates/atuin-client/src/settings.rs
@@ -734,7 +734,13 @@ pub struct Settings {
     pub style: Style,
     pub auto_sync: bool,
     pub update_check: bool,
+
+    /// The address of the Atuin Hub. Used for Hub-specific features like AI.
+    pub hub_address: String,
+
+    /// The sync address for atuin.
     pub sync_address: String,
+
     pub sync_frequency: String,
     pub db_path: String,
     pub record_store_path: String,
@@ -1014,6 +1020,7 @@ impl Settings {
             .set_default("timezone", "local")?
             .set_default("auto_sync", true)?
             .set_default("update_check", cfg!(feature = "check-update"))?
+            .set_default("hub_address", "https://hub.atuin.sh")?
             .set_default("sync_address", "https://api.atuin.sh")?
             .set_default("sync_frequency", "5m")?
             .set_default("search_mode", "fuzzy")?