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

orch/cli/
stream.rs
