1 files changed, 313 insertions, 0 deletions
diff --git a/crates/turtle/src/client/mod.rs b/crates/turtle/src/client/mod.rs
new file mode 100644
index 00000000..531d8597
--- /dev/null
+++ b/crates/turtle/src/client/mod.rs
@@ -0,0 +1,313 @@
+use eyre::{Context as EyreContext, Result};
+use tonic::Code;
+use tonic::transport::{Channel, Endpoint, Uri};
+use tower::service_fn;
+
+use hyper_util::rt::TokioIo;
+
+use tokio::net::UnixStream;
+
+use crate::generated::{
+    self, DAEMON_PROTOCOL_VERSION,
+    control::{
+        ForceSyncReply, ForceSyncRequest, PathsReply, PathsRequest, StatusReply, StatusRequest,
+        control_client::ControlClient as ControlServiceClient,
+    },
+    history::{
+        EndHistoryReply, EndHistoryRequest, HistoryEntry, HistoryRequest, StartHistoryReply,
+        StartHistoryRequest, TailHistoryRequest,
+        history_client::HistoryClient as HistoryServiceClient,
+    },
+};
+
+pub use crate::generated::history::{HistoryEventKind, TailHistoryReply};
+use crate::history::History;
+
+fn normalize_optional_field(value: &str) -> Option<String> {
+    let trimmed = value.trim();
+    if trimmed.is_empty() {
+        None
+    } else {
+        Some(trimmed.to_owned())
+    }
+}
+
+/// The protobuf compile (for some reason) supports not actually sending a request with all fields
+/// (so every field is either fetched from the wire or set to a default).
+/// For custom messages, there are no defaults and thus they get generated as `Option`s.
+/// Our code will (obviously) never leave out a required (!) field in a message, and thus we can
+/// just unwrap all the pointless options.
+fn unwrap_proto_option<T>(field: Option<T>) -> T {
+    field.expect("should be some (see comment above)")
+}
+
+#[must_use]
+pub fn proto_duration_to_std(proto: Option<generated::history::Duration>) -> std::time::Duration {
+    std::time::Duration::from_nanos(unwrap_proto_option(proto).value)
+}
+
+#[must_use]
+pub fn proto_timestamp_to_time(proto: Option<generated::history::Timestamp>) -> OffsetDateTime {
+    OffsetDateTime::from_unix_timestamp_nanos(i128::from(unwrap_proto_option(proto).value))
+        .expect("Daemon history timestamp should always be valid")
+}
+
+#[must_use]
+pub fn time_to_proto_timestamp(t: OffsetDateTime) -> Option<generated::history::Timestamp> {
+    Some(generated::history::Timestamp {
+        value: t.unix_timestamp_nanos() as u64,
+    })
+}
+
+#[must_use]
+pub fn std_to_proto_duration(s: std::time::Duration) -> Option<generated::history::Duration> {
+    Some(generated::history::Duration {
+        value: s.as_nanos() as u64,
+    })
+}
+
+#[must_use]
+pub fn history_entry_to_history(entry: HistoryEntry) -> History {
+    let timestamp = proto_timestamp_to_time(entry.timestamp);
+    let duration = proto_duration_to_std(entry.duration);
+
+    History {
+        id: entry.id.into(),
+        timestamp,
+        duration,
+        exit: entry.exit,
+        command: entry.command,
+        cwd: entry.cwd,
+        session: entry.session,
+        hostname: entry.hostname,
+        author: entry.author,
+        intent: normalize_optional_field(&entry.intent),
+        deleted_at: None,
+    }
+}
+
+#[must_use]
+pub fn daemon_matches_expected(protocol: u32) -> bool {
+    protocol == DAEMON_PROTOCOL_VERSION
+}
+
+#[must_use]
+pub fn daemon_mismatch_message(protocol: u32) -> String {
+    if protocol == DAEMON_PROTOCOL_VERSION {
+        unreachable!()
+    } 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
+}
+
+#[derive(Debug)]
+pub enum Probe {
+    Ready(ControlClient),
+    NeedsRestart(String),
+    Unreachable(eyre::Report),
+}
+
+/// Check if a client can reach the daemon.
+pub async fn probe(path: String) -> Probe {
+    let mut client = match ControlClient::new(path).await {
+        Ok(client) => client,
+        Err(err) => return Probe::Unreachable(err),
+    };
+
+    match client.status().await {
+        Ok(status) => {
+            if daemon_matches_expected(status.protocol) {
+                Probe::Ready(client)
+            } else {
+                Probe::NeedsRestart(daemon_mismatch_message(status.protocol))
+            }
+        }
+        Err(err) => Probe::Unreachable(err),
+    }
+}
+
+// ============================================================================
+// History Client
+// ============================================================================
+
+#[derive(Debug)]
+pub struct HistoryClient {
+    client: HistoryServiceClient<Channel>,
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct Range {
+    pub start: OffsetDateTime,
+    pub end: OffsetDateTime,
+}
+
+pub use time::Duration;
+pub use time::OffsetDateTime;
+
+// 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).max_decoding_message_size({
+            const BIT: usize = 1;
+            const BYTE: usize = 8 * BIT;
+            const KILO_BYTE: usize = 1024 * BYTE;
+            const MEGA_BYTE: usize = 1024 * KILO_BYTE;
+
+            16 * MEGA_BYTE
+        });
+
+        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: time_to_proto_timestamp(h.timestamp),
+            author: h.author,
+            intent: h.intent,
+        };
+
+        Ok(self.client.start_history(req).await?.into_inner())
+    }
+
+    pub async fn end_history(
+        &mut self,
+        id: String,
+        duration: std::time::Duration,
+        exit: i64,
+    ) -> Result<EndHistoryReply> {
+        let req = EndHistoryRequest {
+            id,
+            exit,
+            duration: std_to_proto_duration(duration),
+        };
+
+        Ok(self.client.end_history(req).await?.into_inner())
+    }
+
+    pub async fn history(&mut self, range: Option<Range>) -> Result<Vec<History>> {
+        let req = HistoryRequest {
+            range: range.map(|r| generated::history::Range {
+                start: time_to_proto_timestamp(r.start),
+                end: time_to_proto_timestamp(r.end),
+            }),
+        };
+
+        let reply = self.client.history(req).await?.into_inner();
+
+        Ok(reply
+            .entries
+            .into_iter()
+            .map(history_entry_to_history)
+            .collect())
+    }
+
+    pub 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.
+#[derive(Debug)]
+pub struct ControlClient {
+    client: ControlServiceClient<Channel>,
+}
+
+impl ControlClient {
+    /// Connect to the daemon's control service.
+    pub 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 })
+    }
+
+    pub async fn paths(&mut self) -> Result<PathsReply> {
+        Ok(self.client.paths(PathsRequest {}).await?.into_inner())
+    }
+
+    pub async fn force_sync(&mut self) -> Result<ForceSyncReply> {
+        Ok(self
+            .client
+            .force_sync(ForceSyncRequest {})
+            .await?
+            .into_inner())
+    }
+
+    pub async fn status(&mut self) -> Result<StatusReply> {
+        Ok(self.client.status(StatusRequest {}).await?.into_inner())
+    }
+}