|
|
@@ -137,7 +137,10 @@ async fn wait_ready(base_url: &str) -> Result<()> {
|
|
|
}
|
|
|
|
|
|
/// Tee one output pipe to the console and the log file.
|
|
|
-/// Tee one output pipe to the console and the log file.
|
|
|
+///
|
|
|
+/// The console copy keeps the child's ANSI colors; the file copy is
|
|
|
+/// stripped (state machine below) so the saved .log stays readable in
|
|
|
+/// `sm121 report` and under grep/less.
|
|
|
///
|
|
|
/// The read buffer must not be empty: `read()` on a zero-length slice
|
|
|
/// returns 0 (interpreted as EOF), so an empty `Vec` would silently sever
|
|
|
@@ -145,6 +148,8 @@ async fn wait_ready(base_url: &str) -> Result<()> {
|
|
|
/// (silent exit 120 in CPython).
|
|
|
fn pump<R: Read>(mut reader: R, to_console: bool, file: Arc<Mutex<fs::File>>) {
|
|
|
let mut buf = vec![0u8; 8192];
|
|
|
+ let mut carry = Vec::new(); // trailing partial escape across chunks
|
|
|
+ let mut clean = Vec::new();
|
|
|
loop {
|
|
|
let n = match reader.read(&mut buf) {
|
|
|
Ok(n) => n,
|
|
|
@@ -157,10 +162,71 @@ fn pump<R: Read>(mut reader: R, to_console: bool, file: Arc<Mutex<fs::File>>) {
|
|
|
if to_console {
|
|
|
std::io::stdout().write_all(&buf[..n]).ok();
|
|
|
}
|
|
|
- if let Ok(mut f) = file.lock() {
|
|
|
- f.write_all(&buf[..n]).ok();
|
|
|
+ carry.extend_from_slice(&buf[..n]);
|
|
|
+ strip_ansi_step(&mut carry, &mut clean);
|
|
|
+ if !clean.is_empty() {
|
|
|
+ if let Ok(mut f) = file.lock() {
|
|
|
+ f.write_all(&clean).ok();
|
|
|
+ }
|
|
|
+ clean.clear();
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+/// Strip ANSI escape sequences out of `buf` into `out`, in place. The
|
|
|
+/// consumed prefix is drained from `buf`; a trailing partial escape (a
|
|
|
+/// sequence split across chunk boundaries) is left in `buf` as the carry
|
|
|
+/// for the next call. Handles CSI (`ESC [ ... final`), OSC (`ESC ] ... BEL`
|
|
|
+/// or `ST`), and generic escapes with intermediate bytes (`ESC ( B`).
|
|
|
+fn strip_ansi_step(buf: &mut Vec<u8>, out: &mut Vec<u8>) {
|
|
|
+ let n = buf.len();
|
|
|
+ let mut i = 0;
|
|
|
+ while i < n {
|
|
|
+ if buf[i] != 0x1b {
|
|
|
+ out.push(buf[i]);
|
|
|
+ i += 1;
|
|
|
+ continue;
|
|
|
+ }
|
|
|
+ let rest = &buf[i + 1..];
|
|
|
+ match rest.first() {
|
|
|
+ Some(b'[') => match rest.iter().skip(1).position(|c| (0x40..=0x7e).contains(c)) {
|
|
|
+ Some(p) => i += 1 + 1 + p + 1, // ESC [ params final
|
|
|
+ None => break, // incomplete: wait for more data
|
|
|
+ },
|
|
|
+ Some(b']') => match rest.iter().skip(1).position(|c| *c == 0x07 || *c == 0x1b) {
|
|
|
+ Some(p) => {
|
|
|
+ i += 1 + 1 + p + 1;
|
|
|
+ if rest[1 + p] == 0x1b {
|
|
|
+ i += 1; // consume the backslash of the ST
|
|
|
+ }
|
|
|
+ }
|
|
|
+ None => break,
|
|
|
+ },
|
|
|
+ None => break, // lone ESC at the chunk end
|
|
|
+ Some(_) => {
|
|
|
+ // Generic escape: intermediate bytes (0x20-0x2f), then one
|
|
|
+ // final byte — drop all of them.
|
|
|
+ let mut j = 1;
|
|
|
+ while j < rest.len() && (0x20..=0x2f).contains(&rest[j]) {
|
|
|
+ j += 1;
|
|
|
+ }
|
|
|
+ if j < rest.len() {
|
|
|
+ j += 1; // final byte
|
|
|
+ }
|
|
|
+ i += 1 + j;
|
|
|
+ }
|
|
|
}
|
|
|
}
|
|
|
+ buf.drain(..i);
|
|
|
+}
|
|
|
+
|
|
|
+/// One-shot ANSI strip (whole buffer at once; for files already on disk —
|
|
|
+/// see `report::serve`).
|
|
|
+pub(crate) fn strip_ansi(bytes: &[u8]) -> Vec<u8> {
|
|
|
+ let mut buf = bytes.to_vec();
|
|
|
+ let mut out = Vec::new();
|
|
|
+ strip_ansi_step(&mut buf, &mut out);
|
|
|
+ out
|
|
|
}
|
|
|
|
|
|
/// Run the bench tool in `cwd`, teeing stdout+stderr to the console and the
|
|
|
@@ -400,4 +466,27 @@ mod tests {
|
|
|
r2.args = crate::pre_serve::RecipeArgs::List(vec!["--trust-remote-code".into()]);
|
|
|
assert!(served_model(&r2).is_none());
|
|
|
}
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn strip_ansi_removes_codes_keeps_text() {
|
|
|
+ let raw = b"\x1b[38;5;214mred\x1b[0m plain \x1b]0;title\x07osc \x1b]2;x\x1b\\st \x1b(Bc";
|
|
|
+ assert_eq!(strip_ansi(raw), b"red plain osc st c");
|
|
|
+ // nothing to strip
|
|
|
+ assert_eq!(strip_ansi(b"clean text"), b"clean text".to_vec());
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn strip_ansi_survives_split_chunks() {
|
|
|
+ // the escape sequence straddles the chunk boundary
|
|
|
+ let a = b"ab\x1b[38;5;2"; // split inside a CSI parameter
|
|
|
+ let b2 = b"14mtext";
|
|
|
+ let mut ca = Vec::new();
|
|
|
+ ca.extend_from_slice(a);
|
|
|
+ let mut out = Vec::new();
|
|
|
+ strip_ansi_step(&mut ca, &mut out);
|
|
|
+ ca.extend_from_slice(b2);
|
|
|
+ strip_ansi_step(&mut ca, &mut out);
|
|
|
+ assert_eq!(out, b"abtext");
|
|
|
+ assert!(ca.is_empty(), "no dangling carry: {ca:?}");
|
|
|
+ }
|
|
|
}
|