diff options
Diffstat (limited to 'crates/daemon/src/api/history.rs')
| -rw-r--r-- | crates/daemon/src/api/history.rs | 243 |
1 files changed, 243 insertions, 0 deletions
diff --git a/crates/daemon/src/api/history.rs b/crates/daemon/src/api/history.rs new file mode 100644 index 00000000..078e8c64 --- /dev/null +++ b/crates/daemon/src/api/history.rs @@ -0,0 +1,243 @@ +use std::{pin::Pin, time::Duration}; + +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::store::HistoryStore, settings::Settings}, + daemon::DaemonHandle, + events::DaemonEvent, +}; +use turtle_api::{ + client::{ + proto_duration_to_std, proto_timestamp_to_time, std_to_proto_duration, + time_to_proto_timestamp, + }, + generated::{ + DAEMON_PROTOCOL_VERSION, + history::{ + EndHistoryReply, EndHistoryRequest, HistoryEntry, HistoryEventKind, HistoryReply, + HistoryRequest, StartHistoryReply, StartHistoryRequest, TailHistoryReply, + TailHistoryRequest, + history_server::{History as HistorySvc, HistoryServer}, + }, + }, + history::{History, HistoryId}, +}; + +/// The gRPC service implementation. +/// +/// This is a thin wrapper that delegates to the component's shared state. +pub(crate) struct HistoryService { + /// Commands currently running (not yet completed). + running: DashMap<HistoryId, History>, + + /// Handle to the daemon (set during start). + handle: DaemonHandle, + + /// History store for pushing records + history_store: HistoryStore, +} + +impl HistoryService { + pub(crate) async fn new(handle: DaemonHandle) -> Result<Self> { + let host_id = Settings::host_id().await?; + let history_store = + HistoryStore::new(handle.store().clone(), host_id, *handle.encryption_key()); + + Ok(Self { + running: DashMap::new(), + handle, + history_store, + }) + } + + /// Get a tonic server for this service. + pub(crate) fn into_server(self) -> HistoryServer<Self> { + HistoryServer::new(self) + } +} + +fn history_to_reply(history: History) -> HistoryEntry { + HistoryEntry { + timestamp: time_to_proto_timestamp(history.timestamp), + id: history.id.to_string(), + 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: std_to_proto_duration(history.duration), + } +} + +#[tonic::async_trait] +impl HistorySvc for HistoryService { + type TailHistoryStream = Pin<Box<dyn Stream<Item = Result<TailHistoryReply, Status>> + Send>>; + + #[instrument(skip_all, level = Level::INFO)] + async fn history( + &self, + request: Request<HistoryRequest>, + ) -> Result<Response<HistoryReply>, Status> { + let req = request.into_inner(); + + let entries = if let Some(range) = req.range { + let from = proto_timestamp_to_time(range.start); + let to = proto_timestamp_to_time(range.end); + + self.handle.history_db().range(from, to).await + } else { + self.handle.history_db().list(None, false, false).await + } + .map_err(|e| Status::internal(format!("failed to read db: {e:?}")))? + .into_iter() + .map(history_to_reply) + .collect(); + + Ok(Response::new(HistoryReply { entries })) + } + + #[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 = proto_timestamp_to_time(req.timestamp); + + 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.unwrap_or_default()) + .build() + .into(); + + self.handle.emit(DaemonEvent::HistoryStarted(h.clone())); + + let id = h.id; + tracing::info!(id = id.to_string(), "start history called"); + self.running.insert(id, 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)] + async fn end_history( + &self, + request: Request<EndHistoryRequest>, + ) -> Result<Response<EndHistoryReply>, Status> { + let req = request.into_inner(); + let id = HistoryId::from(req.id); + + tracing::info!(id = id.to_string(), "end history called"); + + if let Some((_, mut history)) = self.running.remove(&id) { + history.exit = req.exit; + history.duration = match proto_duration_to_std(req.duration) { + Duration::ZERO => Duration::from_nanos_u128( + (OffsetDateTime::now_utc() - history.timestamp).whole_nanoseconds() as u128, + ), + value => value, + }; + + self.handle + .history_db() + .save(&history) + .await + .map_err(|e| Status::internal(format!("failed to write to db: {e:?}")))?; + + tracing::info!( + id = id.to_string(), + duration = history.duration.as_nanos(), + "end history" + ); + + let (record_id, idx) = self + .history_store + .push(history.clone()) + .await + .map_err(|e| Status::internal(format!("failed to push record to store: {e:?}")))?; + + self.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)] + async fn tail_history( + &self, + _request: Request<TailHistoryRequest>, + ) -> Result<Response<Self::TailHistoryStream>, Status> { + let mut rx = self.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(TailHistoryReply { + kind: HistoryEventKind::Started.into(), + history: Some(history_to_reply(history)), + }), + DaemonEvent::HistoryEnded(history) => Some(TailHistoryReply { + kind: HistoryEventKind::Ended.into(), + history: Some(history_to_reply(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))) + } +} |