chore: Commit

1 files changed, 180 insertions, 0 deletions

diff --git a/crates/daemon/src/api/client/mod.rs b/crates/daemon/src/api/client/mod.rs
new file mode 100644
index 00000000..d6cbbe85
--- /dev/null
+++ b/crates/daemon/src/api/client/mod.rs
@@ -0,0 +1,180 @@
+use eyre::{Context as EyreContext, Result};
+use tonic::Code;
+use tonic::transport::{Channel, Endpoint, Uri};
+use tower::service_fn;
+
+use hyper_util::rt::TokioIo;
+
+#[cfg(unix)]
+use tokio::net::UnixStream;
+
+use crate::{
+    aclient::{history::History, settings::Settings},
+    api::{
+        DAEMON_PROTOCOL_VERSION, DAEMON_VERSION,
+        generated::{
+            control::{
+                StatusReply, StatusRequest, control_client::ControlClient as ControlServiceClient,
+            },
+            history::{
+                EndHistoryReply, EndHistoryRequest, StartHistoryReply, StartHistoryRequest,
+                TailHistoryReply, TailHistoryRequest,
+                history_client::HistoryClient as HistoryServiceClient,
+            },
+        },
+    },
+};
+
+fn daemon_matches_expected(version: &str, protocol: u32) -> bool {
+    version == DAEMON_VERSION && protocol == DAEMON_PROTOCOL_VERSION
+}
+
+fn daemon_mismatch_message(version: &str, protocol: u32) -> String {
+    if protocol == DAEMON_PROTOCOL_VERSION {
+        format!("daemon is out of date: expected {DAEMON_VERSION}, got {version}")
+    } else {
+        format!("daemon protocol mismatch: expected {DAEMON_PROTOCOL_VERSION}, got {protocol}")
+    }
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum DaemonClientErrorKind {
+    Connect,
+    Unavailable,
+    Unimplemented,
+    Other,
+}
+
+#[must_use]
+pub fn classify_error(error: &eyre::Report) -> DaemonClientErrorKind {
+    for cause in error.chain() {
+        if cause.downcast_ref::<tonic::transport::Error>().is_some() {
+            return DaemonClientErrorKind::Connect;
+        }
+
+        if let Some(status) = cause.downcast_ref::<tonic::Status>() {
+            return match status.code() {
+                Code::Unavailable => DaemonClientErrorKind::Unavailable,
+                Code::Unimplemented => DaemonClientErrorKind::Unimplemented,
+                _ => DaemonClientErrorKind::Other,
+            };
+        }
+    }
+
+    DaemonClientErrorKind::Other
+}
+
+// ============================================================================
+// History Client
+// ============================================================================
+
+pub struct HistoryClient {
+    client: HistoryServiceClient<Channel>,
+}
+
+// Wrap the grpc client
+impl HistoryClient {
+    #[cfg(unix)]
+    pub async fn new(path: String) -> Result<Self> {
+        use eyre::Context;
+
+        let log_path = path.clone();
+        let channel = Endpoint::try_from("http://atuin_local_daemon:0")?
+            .connect_with_connector(service_fn(move |_: Uri| {
+                let path = path.clone();
+
+                async move {
+                    Ok::<_, std::io::Error>(TokioIo::new(UnixStream::connect(path.clone()).await?))
+                }
+            }))
+            .await
+            .wrap_err_with(|| {
+                format!(
+                    "failed to connect to local atuin daemon at {}. Is it running?",
+                    &log_path
+                )
+            })?;
+
+        let client = HistoryServiceClient::new(channel);
+
+        Ok(Self { client })
+    }
+
+    pub async fn start_history(&mut self, h: History) -> Result<StartHistoryReply> {
+        let req = StartHistoryRequest {
+            command: h.command,
+            cwd: h.cwd,
+            hostname: h.hostname,
+            session: h.session,
+            timestamp: h.timestamp.unix_timestamp_nanos() as u64,
+            author: h.author,
+            intent: h.intent.unwrap_or_default(),
+        };
+
+        Ok(self.client.start_history(req).await?.into_inner())
+    }
+
+    pub async fn end_history(
+        &mut self,
+        id: String,
+        duration: u64,
+        exit: i64,
+    ) -> Result<EndHistoryReply> {
+        let req = EndHistoryRequest { id, exit, duration };
+
+        Ok(self.client.end_history(req).await?.into_inner())
+    }
+
+    pub(crate) async fn tail_history(&mut self) -> Result<tonic::Streaming<TailHistoryReply>> {
+        Ok(self
+            .client
+            .tail_history(TailHistoryRequest {})
+            .await?
+            .into_inner())
+    }
+}
+
+// ============================================================================
+// Control Client
+// ============================================================================
+
+/// Client for the Control gRPC service.
+pub struct ControlClient {
+    client: ControlServiceClient<Channel>,
+}
+
+impl ControlClient {
+    /// Connect to the daemon's control service.
+    pub(crate) async fn new(path: String) -> Result<Self> {
+        let log_path = path.clone();
+        let channel = Endpoint::try_from("http://atuin_local_daemon:0")?
+            .connect_with_connector(service_fn(move |_: Uri| {
+                let path = path.clone();
+
+                async move {
+                    Ok::<_, std::io::Error>(TokioIo::new(UnixStream::connect(path.clone()).await?))
+                }
+            }))
+            .await
+            .wrap_err_with(|| {
+                format!(
+                    "failed to connect to local atuin daemon at {}. Is it running?",
+                    &log_path
+                )
+            })?;
+
+        let client = ControlServiceClient::new(channel);
+
+        Ok(Self { client })
+    }
+
+    /// Connect using settings.
+    #[cfg(unix)]
+    pub async fn from_settings(settings: &Settings) -> Result<Self> {
+        Self::new(settings.daemon.socket_path.clone()).await
+    }
+
+    pub async fn status(&mut self) -> Result<StatusReply> {
+        Ok(self.client.status(StatusRequest {}).await?.into_inner())
+    }
+}