orch/cli/
stream.rs

1//! `orch stream` and `orch task watch`: tail each running agent's `output.json` through the
2//! NDJSON formatter. The runner tees agent stdout there, so lines arrive whole
3
4use crate::channels::tmux;
5use crate::cli;
6use crate::cli::ndjson::{self, Formatter};
7use crate::cli::task_ref;
8use crate::home;
9use crate::store::TaskStatus;
10use crate::tmux::TmuxManager;
11use std::collections::HashMap;
12use std::fs::File;
13use std::io::{IsTerminal, Read, Seek, SeekFrom, Write};
14use std::path::{Path, PathBuf};
15use std::time::{Duration, SystemTime};
16
17use tokio::time;
18
19use std::fs;
20use std::io;
21use std::mem;
22
23// Replayed when a file is first opened, so a new viewer sees recent context
24const BACKLOG_BYTES: u64 = 64 * 1024;
25
26struct Tail {
27    path: PathBuf,
28    offset: u64,
29    partial: Vec<u8>,
30    skip_first: bool,
31    formatter: Formatter,
32}
33
34impl Tail {
35    fn open(path: PathBuf, color: bool) -> Self {
36        let len = fs::metadata(&path).map_or(0, |m| m.len());
37        let offset = len.saturating_sub(BACKLOG_BYTES);
38        Self {
39            path,
40            offset,
41            partial: Vec::new(),
42            skip_first: offset > 0,
43            formatter: Formatter::new(color),
44        }
45    }
46
47    /// Complete lines appended since the last call
48    fn read_lines(&mut self) -> Vec<String> {
49        let mut buf = Vec::new();
50        let read = File::open(&self.path).and_then(|mut f| {
51            f.seek(SeekFrom::Start(self.offset))?;
52            f.read_to_end(&mut buf)
53        });
54        let Ok(n) = read else {
55            return vec![];
56        };
57
58        self.offset += n as u64;
59        self.partial.extend_from_slice(&buf);
60        let Some(end) = self.partial.iter().rposition(|&b| b == b'\n') else {
61            return vec![];
62        };
63
64        let rest = self.partial.split_off(end + 1);
65        let complete = mem::replace(&mut self.partial, rest);
66        let mut lines: Vec<String> = String::from_utf8_lossy(&complete)
67            .lines()
68            .map(String::from)
69            .collect();
70        // The backlog starts mid-line
71        if mem::take(&mut self.skip_first) && !lines.is_empty() {
72            lines.remove(0);
73        }
74
75        lines
76    }
77}
78
79/// Follow one task (and its review) when `task_id` is set, else every running agent session
80pub async fn run(task_id: Option<&str>, formatted: bool) -> anyhow::Result<()> {
81    let color = formatted && io::stdout().is_terminal();
82    let watched = if let Some(id) = task_id {
83        println!("Watching task {id}. Ctrl+C to stop");
84        Some(Watched::resolve(id).await?)
85    } else {
86        println!("Watching all running agents. Ctrl+C to stop");
87        None
88    };
89
90    let mut tails: HashMap<String, Tail> = HashMap::new();
91    let mut colors: HashMap<String, &str> = HashMap::new();
92    let mut tag_width = 0;
93    let mut said_waiting = false;
94    loop {
95        let found = match &watched {
96            Some(task) => task.sources(),
97            None => session_sources().await,
98        };
99        if let (Some(task), true, false) = (&watched, found.is_empty(), said_waiting) {
100            println!("{}", task.nothing_yet());
101            said_waiting = true;
102        }
103
104        let live: Vec<String> = found.iter().map(|(key, _)| key.clone()).collect();
105        for (key, path) in found {
106            let tail = tails
107                .entry(key)
108                .or_insert_with(|| Tail::open(path.clone(), color));
109            if tail.path != path {
110                *tail = Tail::open(path, color);
111            }
112        }
113
114        // A session that just ended gets this last read, its final message lands after the work
115        for tail in tails.values_mut() {
116            let mut entries = Vec::new();
117            for line in tail.read_lines() {
118                if formatted {
119                    entries.extend(tail.formatter.format(&line));
120                } else {
121                    entries.push(line);
122                }
123            }
124            if entries.is_empty() {
125                continue;
126            }
127
128            // Output scrolls fast, so every entry carries its task and attempt
129            let tag = tag(&tail.path);
130            tag_width = tag_width.max(tag.chars().count());
131            let next = TAG_COLORS[colors.len() % TAG_COLORS.len()];
132            let code = *colors.entry(key_of(&tail.path)).or_insert(next);
133            let gutter = |show: bool| {
134                let text = if show { tag.as_str() } else { "" };
135                let text = format!("{text:<tag_width$} │");
136                if color {
137                    format!("\x1b[{code}m{text}\x1b[0m")
138                } else {
139                    text
140                }
141            };
142            for line in tag_lines(&entries, formatted, gutter) {
143                println!("{line}");
144            }
145        }
146        tails.retain(|key, _| live.contains(key));
147
148        io::stdout().flush()?;
149        time::sleep(Duration::from_secs(1)).await;
150    }
151}
152
153/// The task `orch task watch <id>` follows, resolved once through the shared CLI helper
154struct Watched {
155    id: String,
156    status: String,
157    /// The task's review dir and its own
158    dirs: Vec<PathBuf>,
159}
160
161impl Watched {
162    async fn resolve(id: &str) -> anyhow::Result<Self> {
163        let store = cli::init_store().await?;
164        // The review dir is watched anyway, so `<id>-review` watches the same pair
165        let id = id.strip_suffix("-review").unwrap_or(id);
166        let task = task_ref::resolve(&store, id).await?;
167        Ok(Self {
168            id: task.display_spec(),
169            status: task.status.clone(),
170            dirs: task_ref::state_dirs(&task)?,
171        })
172    }
173
174    fn sources(&self) -> Vec<(String, PathBuf)> {
175        newest_output(&self.dirs)
176            .map(|path| (self.id.clone(), path))
177            .into_iter()
178            .collect()
179    }
180
181    /// A running task with no output file is a problem, so name the path instead of waiting quietly
182    fn nothing_yet(&self) -> String {
183        if self.status != TaskStatus::InProgress.as_str() {
184            return format!(
185                "No output for task {} yet, waiting for it to start",
186                self.id
187            );
188        }
189
190        let searched: Vec<String> = self
191            .dirs
192            .iter()
193            .map(|dir| dir.join("attempts").display().to_string())
194            .collect();
195        format!(
196            "Task {} is in_progress but no output file was found under {}",
197            self.id,
198            searched.join(" or ")
199        )
200    }
201}
202
203/// Output files of every running agent session, keyed by session name
204async fn session_sources() -> Vec<(String, PathBuf)> {
205    let tmux = TmuxManager::new();
206    tmux::list_orch_sessions()
207        .await
208        .unwrap_or_default()
209        .into_iter()
210        .filter_map(|session| {
211            let (project, task) = tmux.parse_session_name(&session)?;
212            let path = session_output(&project, &task_dir_name(&task))?;
213            Some((session, path))
214        })
215        .collect()
216}
217
218/// Session names replace `:` with `-`, task dirs keep the `internal:` id
219fn task_dir_name(task: &str) -> String {
220    match task.strip_prefix("internal-") {
221        Some(rest) => format!("internal:{rest}"),
222        None => task.to_string(),
223    }
224}
225
226/// A tmux session names the repo without its owner, so look under every owner
227fn session_output(project: &str, name: &str) -> Option<PathBuf> {
228    let state = home::state_dir().ok()?;
229    let dirs: Vec<PathBuf> = subdirs(&state)
230        .iter()
231        .flat_map(|owner| subdirs(owner))
232        .filter(|repo| repo.file_name() == Some(project.as_ref()))
233        .map(|repo| repo.join("tasks").join(name))
234        .collect();
235    newest_output(&dirs)
236}
237
238/// Newest `output.json` under `<dir>/attempts/*` for any of `dirs`
239fn newest_output(dirs: &[PathBuf]) -> Option<PathBuf> {
240    let mut best: Option<(SystemTime, PathBuf)> = None;
241    for dir in dirs {
242        for attempt in subdirs(&dir.join("attempts")) {
243            let output = attempt.join("output.json");
244            let Ok(modified) = output.metadata().and_then(|m| m.modified()) else {
245                continue;
246            };
247            if best.as_ref().is_none_or(|(t, _)| modified > *t) {
248                best = Some((modified, output));
249            }
250        }
251    }
252
253    best.map(|(_, p)| p)
254}
255
256fn subdirs(dir: &Path) -> Vec<PathBuf> {
257    fs::read_dir(dir)
258        .into_iter()
259        .flatten()
260        .flatten()
261        .map(|e| e.path())
262        .filter(|p| p.is_dir())
263        .collect()
264}
265
266const TAG_COLORS: [&str; 6] = ["36", "35", "33", "34", "96", "95"];
267
268/// Entry lines with `gutter(true)` on each step's first line and `gutter(false)` on the rest
269fn tag_lines(entries: &[String], formatted: bool, gutter: impl Fn(bool) -> String) -> Vec<String> {
270    let mut out = Vec::new();
271    for entry in entries {
272        let step = !formatted || !entry.starts_with(ndjson::RESULT_PREFIX);
273        if formatted && step {
274            out.push(gutter(false));
275        }
276
277        for (i, line) in entry.lines().enumerate() {
278            out.push(format!("{} {line}", gutter(step && i == 0)));
279        }
280    }
281    out
282}
283
284/// `<repo>#<task> a<n>` from `.../<repo>/tasks/<task>/attempts/<n>/output.json`
285fn tag(path: &Path) -> String {
286    match path_parts(path).as_slice() {
287        [attempt, _, task, _, repo] => {
288            let task = task.strip_prefix("internal:").unwrap_or(task);
289            format!("{repo}#{task} a{attempt}")
290        }
291        _ => path.display().to_string(),
292    }
293}
294
295/// Task dir of an output file, so a task keeps its colour across attempts
296fn key_of(path: &Path) -> String {
297    path.ancestors()
298        .nth(3)
299        .map_or_else(String::new, |p| p.display().to_string())
300}
301
302fn path_parts(path: &Path) -> Vec<String> {
303    path.iter()
304        .rev()
305        .skip(1)
306        .take(5)
307        .map(|p| p.to_string_lossy().into_owned())
308        .collect()
309}
310
311#[cfg(test)]
312mod tests {
313    use std::fs::OpenOptions;
314    use std::thread;
315
316    use super::*;
317
318    #[test]
319    fn tail_returns_only_complete_lines() {
320        let dir = tempfile::tempdir().unwrap();
321        let path = dir.path().join("output.json");
322        fs::write(&path, "one\ntw").unwrap();
323        let mut tail = Tail::open(path.clone(), false);
324        assert_eq!(tail.read_lines(), vec!["one"]);
325
326        let mut f = OpenOptions::new().append(true).open(&path).unwrap();
327        f.write_all(b"o\nthree\n").unwrap();
328        assert_eq!(tail.read_lines(), vec!["two", "three"]);
329        assert_eq!(tail.read_lines(), [] as [String; 0]);
330    }
331
332    #[test]
333    fn tail_skips_cut_line_in_backlog() {
334        let dir = tempfile::tempdir().unwrap();
335        let path = dir.path().join("output.json");
336        let long = "x".repeat(usize::try_from(BACKLOG_BYTES).unwrap());
337
338        fs::write(&path, format!("{long}\nlast\n")).unwrap();
339        assert_eq!(Tail::open(path, false).read_lines(), vec!["last"]);
340    }
341
342    #[test]
343    fn task_dir_name_restores_internal_colon() {
344        assert_eq!(task_dir_name("internal-8-review"), "internal:8-review");
345        assert_eq!(task_dir_name("3601"), "3601");
346    }
347
348    #[test]
349    fn in_progress_without_output_names_the_path() {
350        let watched = Watched {
351            id: "owner/repo#5".to_string(),
352            status: TaskStatus::InProgress.as_str().to_string(),
353            dirs: vec![
354                PathBuf::from("/state/owner/repo/tasks/5-review"),
355                PathBuf::from("/state/owner/repo/tasks/5"),
356            ],
357        };
358        let msg = watched.nothing_yet();
359        assert!(msg.contains("in_progress"), "{msg}");
360        assert!(msg.contains("/state/owner/repo/tasks/5/attempts"), "{msg}");
361        assert!(
362            msg.contains("/state/owner/repo/tasks/5-review/attempts"),
363            "{msg}"
364        );
365
366        let waiting = Watched {
367            status: TaskStatus::Routed.as_str().to_string(),
368            ..watched
369        };
370        let msg = waiting.nothing_yet();
371        assert!(msg.contains("waiting for it to start"), "{msg}");
372        assert!(!msg.contains("no output file was found"), "{msg}");
373    }
374
375    #[test]
376    fn newest_output_picks_the_latest_attempt() {
377        let dir = tempfile::tempdir().unwrap();
378        let task = dir.path().join("tasks").join("5");
379        for attempt in ["1", "2"] {
380            let d = task.join("attempts").join(attempt);
381            fs::create_dir_all(&d).unwrap();
382            fs::write(d.join("output.json"), attempt).unwrap();
383            // mtime resolution is coarse on some filesystems
384            thread::sleep(Duration::from_millis(10));
385        }
386        let found = newest_output(&[task.join("5-review"), task.clone()]).unwrap();
387        assert_eq!(found, task.join("attempts").join("2").join("output.json"));
388        assert!(newest_output(&[dir.path().join("missing")]).is_none());
389    }
390
391    #[test]
392    fn tag_drops_internal_prefix() {
393        let path = Path::new("/s/owner/bean/tasks/internal:189488-review/attempts/1/output.json");
394        assert_eq!(tag(path), "bean#189488-review a1");
395        assert_eq!(key_of(path), "/s/owner/bean/tasks/internal:189488-review");
396    }
397
398    #[test]
399    fn tag_lines_label_each_step_only() {
400        let entries = vec![
401            "⏺ Read(a.rs)".to_string(),
402            "  ⎿  Read 3 lines".to_string(),
403            "⏺ Fixing it.\n  Then test.".to_string(),
404        ];
405        let gutter = |show: bool| if show { "t1 │" } else { "   │" }.to_string();
406        assert_eq!(
407            tag_lines(&entries, true, gutter),
408            vec![
409                "   │",
410                "t1 │ ⏺ Read(a.rs)",
411                "   │   ⎿  Read 3 lines",
412                "   │",
413                "t1 │ ⏺ Fixing it.",
414                "   │   Then test.",
415            ]
416        );
417    }
418
419    #[test]
420    fn tag_lines_raw_tags_every_line() {
421        let entries = vec!["{\"a\":1}".to_string(), "{\"b\":2}".to_string()];
422        let gutter = |show: bool| if show { "t │" } else { "  │" }.to_string();
423        assert_eq!(
424            tag_lines(&entries, false, gutter),
425            vec!["t │ {\"a\":1}", "t │ {\"b\":2}"]
426        );
427    }
428}
orch/cli/stream.rs

orch/cli/
stream.rs
