batty_cli/shim/
pty_log.rs

1//! PTY log writer: streams raw PTY bytes to a log file so tmux panes can
2//! `tail -f` the output for display.
3//!
4//! Each shim writes to `.batty/shim-logs/<agent-id>.pty.log`. The log is
5//! truncated on shim start and rotated when it exceeds `MAX_LOG_BYTES`.
6
7use std::fs::{self, File, OpenOptions};
8use std::io::{self, Seek, SeekFrom, Write};
9use std::path::{Path, PathBuf};
10
11/// Maximum log size before rotation (50 MB).
12const MAX_LOG_BYTES: u64 = 50 * 1024 * 1024;
13
14/// A writer that appends raw PTY bytes to a log file with size-based rotation.
15pub struct PtyLogWriter {
16    path: PathBuf,
17    file: File,
18    bytes_written: u64,
19}
20
21impl PtyLogWriter {
22    /// Create a new PTY log writer. Truncates any existing log file.
23    pub fn new(path: &Path) -> io::Result<Self> {
24        if let Some(parent) = path.parent() {
25            fs::create_dir_all(parent)?;
26        }
27
28        let file = OpenOptions::new()
29            .create(true)
30            .write(true)
31            .truncate(true)
32            .open(path)?;
33
34        Ok(Self {
35            path: path.to_path_buf(),
36            file,
37            bytes_written: 0,
38        })
39    }
40
41    /// Append raw bytes to the log. Rotates if the file exceeds `MAX_LOG_BYTES`.
42    pub fn write(&mut self, data: &[u8]) -> io::Result<()> {
43        if self.bytes_written + data.len() as u64 > MAX_LOG_BYTES {
44            self.rotate()?;
45        }
46        self.file.write_all(data)?;
47        self.file.flush()?;
48        self.bytes_written += data.len() as u64;
49        Ok(())
50    }
51
52    /// Rotate: truncate the file and reset the counter. Viewers using `tail -F`
53    /// (capital F) will follow the new file automatically.
54    fn rotate(&mut self) -> io::Result<()> {
55        self.file.set_len(0)?;
56        self.file.seek(SeekFrom::Start(0))?;
57        self.bytes_written = 0;
58        Ok(())
59    }
60
61    /// Return the log file path.
62    pub fn path(&self) -> &Path {
63        &self.path
64    }
65}
66
67// ---------------------------------------------------------------------------
68// Tests
69// ---------------------------------------------------------------------------
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use std::io::Read;
75
76    #[test]
77    fn new_creates_file_and_directories() {
78        let tmp = tempfile::tempdir().unwrap();
79        let log_path = tmp.path().join("sub").join("deep").join("agent.pty.log");
80        let writer = PtyLogWriter::new(&log_path).unwrap();
81        assert!(log_path.exists());
82        assert_eq!(writer.path(), log_path);
83    }
84
85    #[test]
86    fn new_truncates_existing_file() {
87        let tmp = tempfile::tempdir().unwrap();
88        let log_path = tmp.path().join("agent.pty.log");
89        fs::write(&log_path, "old content").unwrap();
90
91        let _writer = PtyLogWriter::new(&log_path).unwrap();
92        let content = fs::read_to_string(&log_path).unwrap();
93        assert!(content.is_empty());
94    }
95
96    #[test]
97    fn write_appends_bytes() {
98        let tmp = tempfile::tempdir().unwrap();
99        let log_path = tmp.path().join("agent.pty.log");
100        let mut writer = PtyLogWriter::new(&log_path).unwrap();
101
102        writer.write(b"hello ").unwrap();
103        writer.write(b"world").unwrap();
104
105        let content = fs::read_to_string(&log_path).unwrap();
106        assert_eq!(content, "hello world");
107        assert_eq!(writer.bytes_written, 11);
108    }
109
110    #[test]
111    fn write_preserves_ansi_escapes() {
112        let tmp = tempfile::tempdir().unwrap();
113        let log_path = tmp.path().join("agent.pty.log");
114        let mut writer = PtyLogWriter::new(&log_path).unwrap();
115
116        let ansi = b"\x1b[31mred\x1b[0m normal";
117        writer.write(ansi).unwrap();
118
119        let mut content = Vec::new();
120        File::open(&log_path)
121            .unwrap()
122            .read_to_end(&mut content)
123            .unwrap();
124        assert_eq!(content, ansi);
125    }
126
127    #[test]
128    fn rotate_truncates_and_resets_counter() {
129        let tmp = tempfile::tempdir().unwrap();
130        let log_path = tmp.path().join("agent.pty.log");
131        let mut writer = PtyLogWriter::new(&log_path).unwrap();
132
133        writer.write(b"some data").unwrap();
134        assert!(writer.bytes_written > 0);
135
136        writer.rotate().unwrap();
137        assert_eq!(writer.bytes_written, 0);
138
139        let content = fs::read_to_string(&log_path).unwrap();
140        assert!(content.is_empty());
141    }
142
143    #[test]
144    fn auto_rotation_at_size_limit() {
145        let tmp = tempfile::tempdir().unwrap();
146        let log_path = tmp.path().join("agent.pty.log");
147        let mut writer = PtyLogWriter::new(&log_path).unwrap();
148
149        // Fake near-limit state
150        writer.bytes_written = MAX_LOG_BYTES - 1;
151        writer.write(b"x").unwrap(); // still fits
152        assert_eq!(writer.bytes_written, MAX_LOG_BYTES);
153
154        // Next write triggers rotation
155        writer.write(b"overflow").unwrap();
156        assert_eq!(writer.bytes_written, 8); // "overflow" length after rotation
157
158        let content = fs::read_to_string(&log_path).unwrap();
159        assert_eq!(content, "overflow");
160    }
161
162    #[test]
163    fn path_returns_correct_path() {
164        let tmp = tempfile::tempdir().unwrap();
165        let log_path = tmp.path().join("test.pty.log");
166        let writer = PtyLogWriter::new(&log_path).unwrap();
167        assert_eq!(writer.path(), log_path);
168    }
169
170    #[test]
171    fn empty_write_is_noop() {
172        let tmp = tempfile::tempdir().unwrap();
173        let log_path = tmp.path().join("agent.pty.log");
174        let mut writer = PtyLogWriter::new(&log_path).unwrap();
175        writer.write(b"").unwrap();
176        assert_eq!(writer.bytes_written, 0);
177    }
178}
batty_cli/shim/pty_log.rs

batty_cli/shim/
pty_log.rs
