1 files changed, 85 insertions, 23 deletions

diff --git a/crates/turtle/src/client/mod.rs b/crates/turtle/src/client/mod.rs
index ec97c994..a2322fdd 100644
--- a/crates/turtle/src/client/mod.rs
+++ b/crates/turtle/src/client/mod.rs
@@ -1,14 +1,13 @@
 use eyre::{Context as EyreContext, Result};
-use time::OffsetDateTime;
 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::generated::history::AddHistoryRequest;
 use crate::generated::{
     self, DAEMON_PROTOCOL_VERSION,
     control::{
@@ -34,14 +33,49 @@ fn normalize_optional_field(value: &str) -> Option<String> {
     }
 }
 
+/// 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 = OffsetDateTime::from_unix_timestamp_nanos(i128::from(entry.timestamp))
-        .expect("Daemon history timestamp should always be valid");
+    let timestamp = proto_timestamp_to_time(entry.timestamp);
+    let duration = proto_duration_to_std(entry.duration);
 
     History {
         id: entry.id.into(),
         timestamp,
-        duration: entry.duration,
+        duration,
         exit: entry.exit,
         command: entry.command,
         cwd: entry.cwd,
@@ -135,12 +169,14 @@ pub struct Range {
     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| {
@@ -169,20 +205,57 @@ impl HistoryClient {
             cwd: h.cwd,
             hostname: h.hostname,
             session: h.session,
-            timestamp: h.timestamp.unix_timestamp_nanos() as u64,
+            timestamp: time_to_proto_timestamp(h.timestamp),
             author: h.author,
-            intent: h.intent.unwrap_or_default(),
+            intent: h.intent,
         };
 
         Ok(self.client.start_history(req).await?.into_inner())
     }
 
-    pub async fn history(&mut self, session: String, range: Option<Range>) -> Result<Vec<History>> {
+    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 add_history(
+        &mut self,
+        h: History,
+        duration: std::time::Duration,
+        exit: i64,
+    ) -> Result<EndHistoryReply> {
+        let req = AddHistoryRequest {
+            start: Some(StartHistoryRequest {
+                timestamp: time_to_proto_timestamp(h.timestamp),
+                command: h.command,
+                cwd: h.cwd,
+                session: h.session,
+                hostname: h.hostname,
+                author: h.author,
+                intent: h.intent,
+            }),
+            exit,
+            duration: std_to_proto_duration(duration),
+        };
+
+        Ok(self.client.add_history(req).await?.into_inner())
+    }
+
+    pub async fn history(&mut self, range: Option<Range>) -> Result<Vec<History>> {
         let req = HistoryRequest {
-            session,
             range: range.map(|r| generated::history::Range {
-                start: r.start.unix_timestamp() as u64,
-                end: r.end.unix_timestamp() as u64,
+                start: time_to_proto_timestamp(r.start),
+                end: time_to_proto_timestamp(r.end),
             }),
         };
 
@@ -195,17 +268,6 @@ impl HistoryClient {
             .collect())
     }
 
-    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 async fn tail_history(&mut self) -> Result<tonic::Streaming<TailHistoryReply>> {
         Ok(self
             .client