1 files changed, 25 insertions, 16 deletions

diff --git a/crates/turtle/src/atuin_daemon/components/semantic.rs b/crates/turtle/src/atuin_daemon/components/semantic.rs
index 052c2d73..69ffc134 100644
--- a/crates/turtle/src/atuin_daemon/components/semantic.rs
+++ b/crates/turtle/src/atuin_daemon/components/semantic.rs
@@ -9,6 +9,7 @@ use std::fmt::{Display, Formatter};
 use std::sync::Arc;
 
 use crate::atuin_client::history::{History, HistoryId};
+use crate::atuin_daemon::generated::semantic;
 use eyre::Result;
 use tokio::sync::Mutex;
 use tonic::{Request, Response, Status, Streaming};
@@ -17,7 +18,7 @@ use tracing::{Level, instrument};
 use crate::atuin_daemon::{
     daemon::{Component, DaemonHandle},
     events::DaemonEvent,
-    semantic::{
+    generated::semantic::{
         CommandCapture, CommandOutputReply, CommandOutputRequest, OutputLine, RecordCommandsReply,
         semantic_server::{Semantic as SemanticSvc, SemanticServer},
     },
@@ -244,7 +245,7 @@ impl SemanticState {
     fn command_output_for_ref(
         &self,
         capture_ref: &CaptureRef,
-        ranges: &[crate::atuin_daemon::semantic::OutputRange],
+        ranges: &[semantic::OutputRange],
     ) -> Option<CommandOutputReply> {
         let stored = self
             .sessions
@@ -534,17 +535,14 @@ fn command_output_not_found() -> CommandOutputReply {
     }
 }
 
-fn select_output_ranges(
-    output: &str,
-    ranges: &[crate::atuin_daemon::semantic::OutputRange],
-) -> Vec<OutputLine> {
+fn select_output_ranges(output: &str, ranges: &[semantic::OutputRange]) -> Vec<OutputLine> {
     let lines: Vec<&str> = output.lines().collect();
     if lines.is_empty() {
         return Vec::new();
     }
 
     let ranges = if ranges.is_empty() {
-        vec![crate::atuin_daemon::semantic::OutputRange { start: 0, end: 999 }]
+        vec![semantic::OutputRange { start: 0, end: 999 }]
     } else {
         ranges.to_vec()
     };
@@ -627,9 +625,20 @@ fn log_record(record: &SemanticCommandRecord, message: &'static str) {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use time::OffsetDateTime;
 
+    use crate::{
+        atuin_client::history::{History, HistoryId},
+        atuin_daemon::{
+            components::semantic::{
+                MAX_COMMANDS_PER_SESSION, MAX_SESSIONS, SemanticCommandRecord, SemanticState,
+                SessionCaptures, SessionId, select_output_ranges,
+            },
+            generated::semantic::{self, CommandOutputReply, CommandOutputRequest, OutputLine},
+        },
+        atuin_pty_proxy::CommandCapture,
+    };
+
     fn history(id: &str, session: &str, command: &str) -> History {
         History {
             id: HistoryId(id.to_string()),
@@ -819,8 +828,8 @@ mod tests {
     fn output_ranges_are_line_based_inclusive_and_support_negative_offsets() {
         let output = "zero\none\ntwo\nthree\nfour";
         let ranges = vec![
-            crate::atuin_daemon::semantic::OutputRange { start: 1, end: 2 },
-            crate::atuin_daemon::semantic::OutputRange { start: -2, end: -1 },
+            semantic::OutputRange { start: 1, end: 2 },
+            semantic::OutputRange { start: -2, end: -1 },
         ];
 
         assert_eq!(
@@ -841,8 +850,8 @@ mod tests {
             .collect::<Vec<_>>()
             .join("\n");
         let ranges = vec![
-            crate::atuin_daemon::semantic::OutputRange { start: 0, end: 100 },
-            crate::atuin_daemon::semantic::OutputRange {
+            semantic::OutputRange { start: 0, end: 100 },
+            semantic::OutputRange {
                 start: -100,
                 end: -1,
             },
@@ -859,8 +868,8 @@ mod tests {
     fn output_ranges_can_leave_gaps_for_client_formatting() {
         let output = "zero\none\ntwo\nthree\nfour";
         let ranges = vec![
-            crate::atuin_daemon::semantic::OutputRange { start: 0, end: 1 },
-            crate::atuin_daemon::semantic::OutputRange { start: 4, end: 4 },
+            semantic::OutputRange { start: 0, end: 1 },
+            semantic::OutputRange { start: 4, end: 4 },
         ];
 
         assert_eq!(
@@ -891,8 +900,8 @@ mod tests {
     fn output_ranges_skip_ranges_fully_outside_output() {
         let output = "zero\none\ntwo";
         let ranges = vec![
-            crate::atuin_daemon::semantic::OutputRange { start: 10, end: 20 },
-            crate::atuin_daemon::semantic::OutputRange {
+            semantic::OutputRange { start: 10, end: 20 },
+            semantic::OutputRange {
                 start: -20,
                 end: -10,
             },