southesk 0.0.8

// Copyright 2026 Thomas Axelsson
// SPDX-License-Identifier: MIT

use rmcp::{RoleClient, model::InitializeRequestParams, service::RunningService};
#[cfg(feature = "__dev")]
use {rmcp::transport::CredentialStore, tracing::warn};

use crate::{
    auth_handler::{AuthHandler, BrowserAuth},
    cred_store::{FullCredStore, SharedCredStore},
    error::ClientBuildError,
};

mod connected;
mod disconnected;

/// The Montrose MCP client
///
/// The client must first be connected, then the Montrose API functions can be
/// used. Build a client using [`ClientBuilder`]. The user will by default
/// automatically be requested to authenticate if there is no valid cached OAuth
/// token.
///
/// The [high-level](#high-level-api) API functions provide a Rust-style API,
/// where the type system is used to enforce required and mutually exclusive
/// arguments.
///
/// The [low-level](#low-level-api) API functions provide an API functions that
/// are generated directly from the MCP API. They provide less guardrails and
/// are clunkier to use, but can allow access to functionality not yet provided
/// by the high-level API functions.
///
/// The [raw](#raw-api) API functions allow calling MCP tools with raw JSON
/// data.
///
/// # Examples
/// ```no_run
/// # use southesk::ClientBuilder;
/// # tokio_test::block_on(
/// # async {
/// let montrose = ClientBuilder::new("My Montrose Client").build().await?;
/// let montrose = montrose.connect().await?;
///
/// let accounts = montrose.get_user_accounts().await?;
/// for account in &accounts {
///     dbg!(account);
/// }
///
/// montrose.disconnect().await;
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// # });
/// ```
#[derive(Debug)]
pub struct Client<S: State = Disconnected> {
    name: String,
    auth_handler: Option<Box<dyn AuthHandler>>,
    cred_store: SharedCredStore,
    state: S,
}

mod private {
    pub trait Sealed {}
}
pub trait State: private::Sealed {}

#[doc(hidden)]
#[derive(Debug)]
pub struct Disconnected;

#[doc(hidden)]
#[derive(Debug)]
pub struct Connected {
    client: RunningService<RoleClient, InitializeRequestParams>,
}

impl private::Sealed for Disconnected {}
impl private::Sealed for Connected {}

impl State for Disconnected {}
impl State for Connected {}

const DEFAULT_CRED_USER: &str = "user";

/// The [`Client`] builder.
#[derive(Debug)]
pub struct ClientBuilder {
    client_name: String,
    auth_handler: Option<Box<dyn AuthHandler>>,
    interactive_auth: bool,
    cred_user: Option<String>,
    cred_store: Option<SharedCredStore>,
    #[cfg(feature = "__dev")]
    force_token_refresh: bool,
}

impl ClientBuilder {
    /// Creates a new builder for [`Client`].
    ///
    /// `client_name` is used to identify the client towards the MCP service. It
    /// is recommended to name it after your application.
    pub fn new(client_name: impl Into<String>) -> Self {
        Self {
            client_name: client_name.into(),
            auth_handler: None,
            interactive_auth: true,
            cred_user: None,
            cred_store: None,
            #[cfg(feature = "__dev")]
            force_token_refresh: false,
        }
    }

    /// Overrides how the user is requested to authenticate.
    ///
    /// [`BrowserAuth`] is used by default.
    #[must_use]
    pub fn auth_handler(mut self, handler: impl AuthHandler + 'static) -> Self {
        self.auth_handler = Some(Box::new(handler));
        self
    }

    /// Disables initiating new interactive authentications (through
    /// [`AuthHandler`]).
    ///
    /// The client will only rely on existing OAuth tokens (from a previous
    /// session). This can be useful for long-running background programs.
    /// [`Client::connect`] will fail if the server denies the existing OAuth
    /// credentials.
    #[must_use]
    pub fn no_auth(mut self) -> Self {
        self.interactive_auth = false;
        self
    }

    /// Sets the user identifier used for the default credential storage.
    ///
    /// This option can be used if the current computer user needs to store
    /// credentials for multiple Montrose accounts or sessions (e.g. for
    /// testing).
    ///
    /// This option is not valid if a custom credential store is provided with
    /// [`ClientBuilder::cred_store`].
    #[must_use]
    pub fn cred_user(mut self, user: impl Into<String>) -> Self {
        self.cred_user = Some(user.into());
        self
    }

    /// Overrides the credential store used to store the user's OAuth credentials
    #[must_use]
    pub fn cred_store(mut self, cred_store: impl FullCredStore + 'static) -> Self {
        self.cred_store = Some(SharedCredStore::new(cred_store));
        self
    }

    #[cfg(feature = "__dev")]
    /// Makes the client force an OAuth token refresh when connecting.
    #[must_use]
    pub fn dev_force_token_refresh(mut self) -> Self {
        self.force_token_refresh = true;
        self
    }

    /// Builds the client.
    ///
    /// The client needs to be connected before the MCP API can be used.
    pub async fn build(self) -> Result<Client<Disconnected>, ClientBuildError> {
        let auth_handler = if self.interactive_auth {
            Some(match self.auth_handler {
                Some(handler) => handler,
                None => Box::new(BrowserAuth::new().await.map_err(|e| ClientBuildError {
                    msg: e.to_string(),
                    source: Some(Box::new(e)),
                })?),
            })
        } else {
            // Having no authentication handler is better than having one that
            // instantly returns an error, as there are some communication with
            // the authentication server before the authentication handler is
            // called.
            None
        };

        let cred_store = if let Some(store) = self.cred_store {
            if self.cred_user.is_some() {
                return Err(ClientBuildError {
                    msg: "cannot specify both a custom credential store and a credential user"
                        .to_string(),
                    source: None,
                });
            }
            store
        } else {
            let cred_user = self.cred_user.unwrap_or(DEFAULT_CRED_USER.to_string());
            SharedCredStore::new({
                #[cfg(feature = "keyring")]
                {
                    use crate::cred_store::KeyringCredStore;

                    KeyringCredStore::new(&self.client_name, &cred_user).map_err(|e| {
                        ClientBuildError {
                            msg: "failed to create keyring credential store".to_string(),
                            source: Some(Box::new(e)),
                        }
                    })?
                }
                #[cfg(not(feature = "keyring"))]
                {
                    use crate::cred_store::PlaintextCredStore;
                    use etcetera::AppStrategy;

                    let client_dirs = etcetera::choose_app_strategy(etcetera::AppStrategyArgs {
                        top_level_domain: String::new(),
                        author: String::new(),
                        app_name: self.client_name.clone(),
                    })
                    .map_err(|e| ClientBuildError {
                        msg: e.to_string(),
                        source: Some(Box::new(e)),
                    })?;
                    PlaintextCredStore::new(
                        client_dirs
                            .state_dir()
                            .unwrap_or_else(|| client_dirs.data_dir())
                            // TOOD: Sanitise the username. Right now, the
                            // code will fail with an error if the username
                            // cannot be used in a filename.
                            .join(format!("{cred_user}_credentials.json")),
                    )
                }
            })
        };

        #[cfg(feature = "__dev")]
        if self.force_token_refresh {
            warn!("Force token refresh enabled");
            let creds = cred_store.load().await.map_err(|e| ClientBuildError {
                msg: "failed to load credentials for setting forced refresh".to_string(),
                source: Some(Box::new(e)),
            })?;
            if let Some(mut creds) = creds {
                use std::time::Duration;

                creds
                    .token_response
                    .as_mut()
                    .ok_or_else(|| ClientBuildError {
                        msg: "no token response in stored credentials".to_string(),
                        source: None,
                    })?
                    .set_expires_in(Some(&Duration::ZERO));
                cred_store.save(creds).await.map_err(|e| ClientBuildError {
                    msg: "failed to save credentials for setting forced refresh".to_string(),
                    source: Some(Box::new(e)),
                })?;
            }
        }

        Ok(Client {
            name: self.client_name,
            auth_handler,
            cred_store,
            state: Disconnected,
        })
    }
}