diff options
Diffstat (limited to 'crates/daemon/src/api')
| -rw-r--r-- | crates/daemon/src/api/client/mod.rs | 180 | ||||
| -rw-r--r-- | crates/daemon/src/api/generated.rs | 28 | ||||
| -rw-r--r-- | crates/daemon/src/api/mod.rs | 6 | ||||
| -rw-r--r-- | crates/daemon/src/api/server/control.rs | 61 | ||||
| -rw-r--r-- | crates/daemon/src/api/server/history.rs | 229 | ||||
| -rw-r--r-- | crates/daemon/src/api/server/mod.rs | 2 |
6 files changed, 506 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()) + } +} diff --git a/crates/daemon/src/api/generated.rs b/crates/daemon/src/api/generated.rs new file mode 100644 index 00000000..304edcd9 --- /dev/null +++ b/crates/daemon/src/api/generated.rs @@ -0,0 +1,28 @@ +#![expect( + unreachable_pub, + unused_qualifications, + clippy::doc_markdown, + clippy::default_trait_access, + clippy::too_many_lines, + clippy::trivially_copy_pass_by_ref, + clippy::allow_attributes, + clippy::derive_partial_eq_without_eq, + reason = "All of these lints are triggered by the generated code" +)] + +/// Semantic command capture gRPC service types. +pub(crate) mod semantic { + tonic::include_proto!("semantic"); +} + +/// History module for the daemon gRPC history service. +/// +/// This module contains the proto-generated types for the history gRPC service. +pub(crate) mod history { + tonic::include_proto!("history"); +} + +/// Control module for external control. +pub(crate) mod control { + tonic::include_proto!("control"); +} diff --git a/crates/daemon/src/api/mod.rs b/crates/daemon/src/api/mod.rs new file mode 100644 index 00000000..e6f8f1f0 --- /dev/null +++ b/crates/daemon/src/api/mod.rs @@ -0,0 +1,6 @@ +pub(crate) mod client; +pub(crate) mod server; +pub(crate) mod generated; + +const DAEMON_VERSION: &str = env!("CARGO_PKG_VERSION"); +const DAEMON_PROTOCOL_VERSION: u32 = 1; diff --git a/crates/daemon/src/api/server/control.rs b/crates/daemon/src/api/server/control.rs new file mode 100644 index 00000000..63fef340 --- /dev/null +++ b/crates/daemon/src/api/server/control.rs @@ -0,0 +1,61 @@ +use tonic::{Request, Response, Status}; +use tracing::{Level, instrument}; + +use crate::{ + api::{ + DAEMON_PROTOCOL_VERSION, DAEMON_VERSION, + generated::control::{ + ForceSyncReply, ForceSyncRequest, StatusReply, StatusRequest, + control_server::{Control, ControlServer}, + }, + }, + daemon::DaemonHandle, +}; + +/// The Control gRPC service. +/// +/// This service is used by external processes to inject events into the daemon. +/// It's not a component - it's part of the daemon's core infrastructure. +pub(crate) struct ControlService { + handle: DaemonHandle, +} + +impl ControlService { + /// Create a new control service with the given daemon handle. + pub(crate) fn new(handle: DaemonHandle) -> Self { + Self { handle } + } + + /// Get a tonic server for this service. + pub(crate) fn into_server(self) -> ControlServer<Self> { + ControlServer::new(self) + } +} + +#[tonic::async_trait] +impl Control for ControlService { + #[instrument(skip_all, level = Level::INFO)] + async fn status( + &self, + _request: Request<StatusRequest>, + ) -> Result<Response<StatusReply>, Status> { + let reply = StatusReply { + healthy: true, + version: DAEMON_VERSION.to_owned(), + pid: std::process::id(), + protocol: DAEMON_PROTOCOL_VERSION, + }; + + Ok(Response::new(reply)) + } + + #[instrument(skip_all, level = Level::INFO)] + async fn force_sync( + &self, + _request: Request<ForceSyncRequest>, + ) -> Result<Response<ForceSyncReply>, Status> { + let reply = ForceSyncReply { accepted: false }; + + Ok(Response::new(reply)) + } +} diff --git a/crates/daemon/src/api/server/history.rs b/crates/daemon/src/api/server/history.rs new file mode 100644 index 00000000..032876a6 --- /dev/null +++ b/crates/daemon/src/api/server/history.rs @@ -0,0 +1,229 @@ +use std::pin::Pin; + +use dashmap::DashMap; +use eyre::Result; +use time::OffsetDateTime; +use tokio_stream::Stream; +use tonic::{Request, Response, Status}; +use tracing::{Level, instrument}; + +use crate::{ + aclient::history::{History, HistoryId, store::HistoryStore}, + api::{ + DAEMON_PROTOCOL_VERSION, + generated::history::{ + EndHistoryReply, EndHistoryRequest, HistoryEntry, HistoryEventKind, StartHistoryReply, + StartHistoryRequest, TailHistoryReply, TailHistoryRequest, + history_server::History as HistorySvc, + }, + }, + daemon::DaemonHandle, + events::DaemonEvent, +}; + +/// The gRPC service implementation. +/// +/// This is a thin wrapper that delegates to the component's shared state. +pub(crate) struct HistoryGrpcService { + /// Commands currently running (not yet completed). + running: DashMap<HistoryId, History>, + + /// Handle to the daemon (set during start). + pub(crate) handle: tokio::sync::RwLock<Option<DaemonHandle>>, + + /// History store for pushing records (set during start). + pub(crate) history_store: tokio::sync::RwLock<Option<HistoryStore>>, +} + +impl HistoryGrpcService { + /// Create a new history component. + pub(crate) fn new() -> Self { + Self { + running: DashMap::new(), + handle: tokio::sync::RwLock::new(None), + history_store: tokio::sync::RwLock::new(None), + } + } +} + +fn history_to_tail_reply(kind: HistoryEventKind, history: History) -> TailHistoryReply { + TailHistoryReply { + kind: kind as i32, + history: Some(HistoryEntry { + timestamp: history.timestamp.unix_timestamp_nanos() as u64, + id: history.id.0, + command: history.command, + cwd: history.cwd, + session: history.session, + hostname: history.hostname, + author: history.author, + intent: history.intent.unwrap_or_default(), + exit: history.exit, + duration: history.duration, + }), + } +} + +#[tonic::async_trait] +impl HistorySvc for HistoryGrpcService { + type TailHistoryStream = Pin<Box<dyn Stream<Item = Result<TailHistoryReply, Status>> + Send>>; + + #[instrument(skip_all, level = Level::INFO)] + async fn start_history( + &self, + request: Request<StartHistoryRequest>, + ) -> Result<Response<StartHistoryReply>, Status> { + let req = request.into_inner(); + + let timestamp = OffsetDateTime::from_unix_timestamp_nanos(i128::from(req.timestamp)) + .map_err(|_| { + Status::invalid_argument( + "failed to parse timestamp as unix time (expected nanos since epoch)", + ) + })?; + + let h: History = History::daemon() + .timestamp(timestamp) + .command(req.command) + .cwd(req.cwd) + .session(req.session) + .hostname(req.hostname) + .author(req.author) + .intent(req.intent) + .build() + .into(); + + // Emit the event + if let Some(handle) = self.handle.read().await.as_ref() { + handle.emit(DaemonEvent::HistoryStarted(h.clone())); + } + + let id = h.id.clone(); + tracing::info!(id = id.to_string(), "start history"); + self.running.insert(id.clone(), h); + + let reply = StartHistoryReply { + id: id.to_string(), + version: env!("CARGO_PKG_VERSION").to_string(), + protocol: DAEMON_PROTOCOL_VERSION, + }; + + Ok(Response::new(reply)) + } + + #[instrument(skip_all, level = Level::INFO)] + #[expect(clippy::significant_drop_tightening, reason = "Would be a logic-bug")] + async fn end_history( + &self, + request: Request<EndHistoryRequest>, + ) -> Result<Response<EndHistoryReply>, Status> { + let req = request.into_inner(); + let id = HistoryId(req.id); + + if let Some((_, mut history)) = self.running.remove(&id) { + history.exit = req.exit; + history.duration = match req.duration { + 0 => i64::try_from( + (OffsetDateTime::now_utc() - history.timestamp).whole_nanoseconds(), + ) + .expect("failed to convert calculated duration to i64"), + value => i64::try_from(value).expect("failed to get i64 duration"), + }; + + // Get the handle and store to save the history + let handle_guard = self.handle.read().await; + let handle = handle_guard + .as_ref() + .ok_or_else(|| Status::internal("component not initialized"))?; + + let store_guard = self.history_store.read().await; + let history_store = store_guard + .as_ref() + .ok_or_else(|| Status::internal("component not initialized"))?; + + // Save to database + handle + .history_db() + .save(&history) + .await + .map_err(|e| Status::internal(format!("failed to write to db: {e:?}")))?; + + tracing::info!(id = id.0, duration = history.duration, "end history"); + + // Push to record store + let (record_id, idx) = history_store + .push(history.clone()) + .await + .map_err(|e| Status::internal(format!("failed to push record to store: {e:?}")))?; + + // Emit the event + handle.emit(DaemonEvent::HistoryEnded(history)); + + let reply = EndHistoryReply { + id: record_id.0.to_string(), + idx, + version: env!("CARGO_PKG_VERSION").to_string(), + protocol: DAEMON_PROTOCOL_VERSION, + }; + + return Ok(Response::new(reply)); + } + + Err(Status::not_found(format!( + "could not find history with id: {id}" + ))) + } + + #[instrument(skip_all, level = Level::INFO)] + #[expect(clippy::significant_drop_tightening, reason = "Would be a logic-bug")] + async fn tail_history( + &self, + _request: Request<TailHistoryRequest>, + ) -> Result<Response<Self::TailHistoryStream>, Status> { + let handle_guard = self.handle.read().await; + let handle = handle_guard + .as_ref() + .cloned() + .ok_or_else(|| Status::internal("component not initialized"))?; + + let mut rx = handle.subscribe(); + let (tx, out_rx) = tokio::sync::mpsc::channel::<Result<TailHistoryReply, Status>>(128); + + tokio::spawn(async move { + loop { + let event = match rx.recv().await { + Ok(event) => event, + Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => { + drop( + tx.send(Err(Status::resource_exhausted(format!( + "tail stream lagged behind and dropped {skipped} events" + )))) + .await, + ); + break; + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + }; + + let reply = match event { + DaemonEvent::HistoryStarted(history) => { + Some(history_to_tail_reply(HistoryEventKind::Started, history)) + } + DaemonEvent::HistoryEnded(history) => { + Some(history_to_tail_reply(HistoryEventKind::Ended, history)) + } + _ => None, + }; + + if let Some(reply) = reply + && tx.send(Ok(reply)).await.is_err() + { + break; + } + } + }); + + let stream = tokio_stream::wrappers::ReceiverStream::new(out_rx); + Ok(Response::new(Box::pin(stream))) + } +} diff --git a/crates/daemon/src/api/server/mod.rs b/crates/daemon/src/api/server/mod.rs new file mode 100644 index 00000000..8d475fe9 --- /dev/null +++ b/crates/daemon/src/api/server/mod.rs @@ -0,0 +1,2 @@ +pub(crate) mod control; +pub(crate) mod history; |
