magi-code 0.96.1

Repository-aware CLI coding agent for terminal work
Documentation
use super::*;

#[derive(Debug, Default)]
pub(crate) struct StreamingControlSanitizer {
    pub(super) state: ControlState,
}

#[derive(Debug, Default)]
pub(super) enum ControlState {
    #[default]
    Normal,
    Escape,
    Csi,
    String {
        escape: bool,
    },
}

impl StreamingControlSanitizer {
    pub(crate) fn push(&mut self, text: &str) -> String {
        let mut out = String::new();
        for ch in text.chars() {
            match &mut self.state {
                ControlState::Normal => {
                    if ch == '\u{1b}' {
                        self.state = ControlState::Escape;
                    } else if !ch.is_control() || matches!(ch, '\n' | '\t' | '\r') {
                        out.push(ch);
                    }
                }
                ControlState::Escape => {
                    self.state = match ch {
                        '[' => ControlState::Csi,
                        ']' | 'P' | '^' | '_' | 'X' => ControlState::String { escape: false },
                        _ => ControlState::Normal,
                    };
                }
                ControlState::Csi => {
                    if ('@'..='~').contains(&ch) {
                        self.state = ControlState::Normal;
                    }
                }
                ControlState::String { escape } => {
                    if ch == '\u{7}' {
                        self.state = ControlState::Normal;
                    } else if *escape {
                        self.state = if ch == '\\' {
                            ControlState::Normal
                        } else if ch == '\u{1b}' {
                            ControlState::String { escape: true }
                        } else {
                            ControlState::String { escape: false }
                        };
                    } else if ch == '\u{1b}' {
                        *escape = true;
                    }
                }
            }
        }
        out
    }

    pub(crate) fn reset(&mut self) {
        self.state = ControlState::Normal;
    }
}

pub(crate) fn sanitize_preview(text: &str) -> String {
    redact_sensitive_text(strip_control_sequences(text).as_str())
}

fn strip_control_sequences(text: &str) -> String {
    let mut out = String::new();
    let mut chars = text.chars().peekable();
    while let Some(ch) = chars.next() {
        if ch != '\u{1b}' {
            if ch.is_control() && !matches!(ch, '\n' | '\t' | '\r') {
                continue;
            }
            out.push(ch);
            continue;
        }

        match chars.peek().copied() {
            Some('[') => {
                chars.next();
                consume_csi(&mut chars);
            }
            Some(']') | Some('P') | Some('^') | Some('_') | Some('X') => {
                chars.next();
                consume_string_control(&mut chars);
            }
            Some(ch) if ('@'..='_').contains(&ch) => {
                chars.next();
            }
            Some(_) => {
                chars.next();
            }
            None => {}
        }
    }
    out
}

fn consume_csi(chars: &mut std::iter::Peekable<std::str::Chars<'_>>) {
    for ch in chars.by_ref() {
        if ('@'..='~').contains(&ch) {
            break;
        }
    }
}

fn consume_string_control(chars: &mut std::iter::Peekable<std::str::Chars<'_>>) {
    while let Some(ch) = chars.next() {
        if ch == '\u{7}' {
            break;
        }
        if ch == '\u{1b}' && chars.peek().copied() == Some('\\') {
            chars.next();
            break;
        }
    }
}