bashkit 0.20.1

//! LLM tool wrapper around a [`Terminal`] session.
//!
//! Decisions:
//! - Not an implementation of the crate's `Tool` trait. That contract is one
//!   isolated shell per call (`{"commands": ...}` in, stdout/stderr out). A
//!   terminal is a session: state, open programs and the screen carry over
//!   between calls. Forcing it into `ToolRequest` would hide that, so this
//!   type has its own schema and keeps the same metadata names
//!   (`name`, `description`, `system_prompt`, `input_schema`,
//!   `output_schema`, `tool_definition`).
//! - Keys use Vim notation inside one string (`ihello<Esc>:wq<Enter>`). LLMs
//!   already know it, it keeps ordering between text and special keys, and it
//!   avoids raw control characters in JSON.
//! - Each call waits for the session to go idle, bounded by `wait_ms`
//!   (default 5 s, max 60 s). A command still running is reported, not
//!   killed; the next call (even with no input) keeps waiting, and `<C-c>`
//!   interrupts it.
//! - `wait_for` (regex) returns as soon as matching text appears in command
//!   output printed during the call. It matches command stdout/stderr only,
//!   never the echoed keystrokes or prompts, so `echo ready<Enter>` with
//!   `wait_for: "ready"` waits for the output line, not the typed command.
//!   Polls in 50 ms slices; the session keeps running between slices.
//! - `screen: "changes"` returns only rows that differ from the previous
//!   call (any mode updates the baseline), to keep long sessions cheap.
//! - Multiple sessions (like terminal tabs) need [`TerminalTool::with_sessions`],
//!   which takes a builder factory because `BashBuilder` is not `Clone`.
//!   Extra sessions share the main session's filesystem, so a file written in
//!   one is visible in the others; shell state (variables, cwd, jobs) is per
//!   session. While a call waits on one session the others keep running, so a
//!   `watch` or long loop in one tab progresses while the agent works in
//!   another. At most 8 sessions; `close` drops one, and an exited session
//!   restarts on the next call that types into it.

use std::time::Duration;

use serde_json::{Value, json};

use super::{CommandRecord, Terminal, TerminalActivity, TerminalSize, TerminalStatus};
use crate::BashBuilder;

/// Tool name used in schemas and definitions.
pub const TERMINAL_TOOL_NAME: &str = "terminal";
const DEFAULT_WAIT_MS: u64 = 5_000;
const MAX_WAIT_MS: u64 = 60_000;
/// Longest `input` string accepted per call.
// THREAT[TM-DOS-119]: tool input feeds the bounded terminal queue; reject
// oversized calls up front instead of silently truncating.
const MAX_INPUT_BYTES: usize = 64 * 1024;
/// Longest `wait_for` pattern, and the compiled-regex size cap.
// THREAT[TM-DOS-119]: the pattern is model-supplied; bound compile cost.
const MAX_PATTERN_BYTES: usize = 1024;
const MAX_REGEX_SIZE: usize = 1 << 20;
const WAIT_SLICE: Duration = Duration::from_millis(50);
/// Name of the session every tool starts with.
pub const MAIN_SESSION: &str = "main";
// THREAT[TM-DOS-119]: each session is a full interpreter; cap them.
const MAX_SESSIONS: usize = 8;
const MAX_SESSION_NAME: usize = 32;

/// Builds a fresh [`BashBuilder`] for each session of a multi-session tool.
pub type SessionFactory = Box<dyn Fn() -> BashBuilder + Send + Sync>;

const DESCRIPTION: &str = "Interactive bash terminal in a sandbox. Type keys, get the screen back. \
Use it for full-screen programs such as vi and less, or a shell session that persists between calls.";

/// Error from [`TerminalTool::call`]. The message is safe to show to the LLM.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("{0}")]
pub struct TerminalToolError(String);

/// A persistent terminal session exposed as an LLM tool.
///
/// One instance is one session: keep it for the whole conversation.
///
/// ```rust
/// use bashkit::Bash;
/// use bashkit::terminal::TerminalTool;
/// use serde_json::json;
///
/// # #[tokio::main]
/// # async fn main() {
/// let mut tool = TerminalTool::new(Bash::builder());
/// let out = tool
///     .call(json!({"input": "echo hi<Enter>"}))
///     .await
///     .unwrap();
/// assert_eq!(out["activity"], "prompt");
/// assert_eq!(out["commands"][0]["output"], "hi\n");
/// assert!(out["screen"].as_str().unwrap().contains("$ echo hi\nhi"));
/// # }
/// ```
pub struct TerminalTool {
    /// Index 0 is the main session.
    sessions: Vec<ToolSession>,
    factory: Option<SessionFactory>,
    size: TerminalSize,
}

struct ToolSession {
    name: String,
    terminal: Terminal,
    /// Screen rows at the previous report, for `screen: "changes"`.
    last_rows: Option<Vec<String>>,
}

impl ToolSession {
    fn new(name: &str, terminal: Terminal) -> Self {
        Self {
            name: name.to_string(),
            terminal,
            last_rows: None,
        }
    }
}

/// How much of the screen a call returns.
#[derive(Clone, Copy, PartialEq, Eq)]
enum ScreenMode {
    Full,
    Changes,
    None,
}

impl TerminalTool {
    /// Start a session on an 80x24 terminal.
    pub fn new(builder: BashBuilder) -> Self {
        Self::with_size(builder, TerminalSize::default())
    }

    /// Start a session on a terminal of the given size.
    pub fn with_size(builder: BashBuilder, size: TerminalSize) -> Self {
        Self {
            sessions: vec![ToolSession::new(
                MAIN_SESSION,
                Terminal::with_size(builder, size),
            )],
            factory: None,
            size,
        }
    }

    /// A tool that can open more sessions on demand (the `session`
    /// argument). `factory` builds the shell for each one; sessions after
    /// the first share its filesystem.
    ///
    /// ```rust
    /// use bashkit::Bash;
    /// use bashkit::terminal::{TerminalSize, TerminalTool};
    /// use serde_json::json;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let mut tool = TerminalTool::with_sessions(TerminalSize::default(), Box::new(Bash::builder));
    /// tool.call(json!({"input": "echo hi > /tmp/f<Enter>"})).await.unwrap();
    /// let out = tool
    ///     .call(json!({"session": "logs", "input": "cat /tmp/f<Enter>"}))
    ///     .await
    ///     .unwrap();
    /// assert_eq!(out["commands"][0]["output"], "hi\n");
    /// # }
    /// ```
    pub fn with_sessions(size: TerminalSize, factory: SessionFactory) -> Self {
        let main = Terminal::with_size(factory(), size);
        Self {
            sessions: vec![ToolSession::new(MAIN_SESSION, main)],
            factory: Some(factory),
            size,
        }
    }

    /// The main session's terminal, for example to read files with
    /// [`Terminal::fs`] or render [`Terminal::take_output`].
    pub fn terminal(&self) -> &Terminal {
        &self.sessions[0].terminal
    }

    /// Mutable access to the main session's terminal, to drive it directly
    /// between tool calls.
    pub fn terminal_mut(&mut self) -> &mut Terminal {
        &mut self.sessions[0].terminal
    }

    /// The terminal of a named session, if it is open.
    pub fn session(&self, name: &str) -> Option<&Terminal> {
        self.sessions
            .iter()
            .find(|s| s.name == name)
            .map(|s| &s.terminal)
    }

    /// Names of the open sessions, main first.
    pub fn session_names(&self) -> Vec<String> {
        self.sessions.iter().map(|s| s.name.clone()).collect()
    }

    fn new_terminal(&self) -> Option<Terminal> {
        let factory = self.factory.as_ref()?;
        let fs = self.sessions[0].terminal.fs();
        Some(Terminal::with_size(factory().fs(fs), self.size))
    }

    /// Tool name: `terminal`.
    pub fn name(&self) -> &str {
        TERMINAL_TOOL_NAME
    }

    /// One-sentence description for tool listings.
    pub fn description(&self) -> &str {
        DESCRIPTION
    }

    /// Terse usage guide for a system prompt.
    pub fn system_prompt(&self) -> String {
        "terminal: interactive sandboxed bash session that persists between calls. \
Send keys in `input` using Vim notation: text is typed as-is, <Enter> runs a line, \
<Esc>, <Tab>, <BS>, <Up>/<Down>/<Left>/<Right>, <PageUp>/<PageDown>, <C-c> interrupts, \
<C-d> ends input, <lt> types a literal '<'. Each call returns the screen, the activity \
(prompt, continuation, running, input, exited) and every command that finished with its exact \
output and exit code. vi and less work; quit vi with <Esc>:wq<Enter>, less with q. \
If activity is input, a command is asking a question (input_prompt holds it); type the answer \
and <Enter>. If activity is running, call again (input may be empty) to keep waiting, or send <C-c>. \
Set wait_for to a regex to return as soon as command output matches (for example a server's \
'listening' line); matched tells whether it did. Set screen to \"changes\" to get only the rows \
that changed since the last call, or \"none\" to skip the screen."
            .to_string()
            + if self.factory.is_some() {
                " Set session to a name to use another terminal tab (created on first use, files \
shared, shell state separate; other tabs keep running); close: true closes it."
            } else {
                ""
            }
    }

    /// JSON Schema for [`call`](Self::call) arguments.
    pub fn input_schema(&self) -> Value {
        let mut schema = json!({
            "type": "object",
            "properties": {
                "input": {
                    "type": "string",
                    "description": "Keys to type, Vim notation: `ls -la<Enter>`, `ihello<Esc>:wq<Enter>`, `<C-c>`. May be empty to just wait and look."
                },
                "wait_ms": {
                    "type": "integer",
                    "minimum": 0,
                    "maximum": MAX_WAIT_MS,
                    "description": "How long to wait for the session to need input (default 5000)."
                },
                "wait_for": {
                    "type": "string",
                    "description": "Regex (multi-line: ^ and $ match at line ends). Return as soon as command output printed during this call matches it, even if the command is still running. Typed keys and prompts are not matched."
                },
                "screen": {
                    "type": "string",
                    "enum": ["full", "changes", "none"],
                    "description": "full (default): whole screen. changes: only rows that differ from the previous call, as screen_changes. none: no screen."
                }
            }
        });
        if self.factory.is_some() {
            schema["properties"]["session"] = json!({
                "type": "string",
                "maxLength": MAX_SESSION_NAME,
                "description": "Terminal tab to use (default \"main\"). A new name opens a new tab sharing the same files; up to 8 tabs."
            });
            schema["properties"]["close"] = json!({
                "type": "boolean",
                "description": "Close the named session (not main) instead of typing into it."
            });
        }
        schema
    }

    /// JSON Schema for the value [`call`](Self::call) returns.
    pub fn output_schema(&self) -> Value {
        json!({
            "type": "object",
            "required": ["activity", "commands", "waiting_for_input"],
            "properties": {
                "screen": {"type": "string", "description": "Visible screen as plain text (screen = full)"},
                "screen_changes": {
                    "type": "array",
                    "description": "Rows that changed since the previous call (screen = changes)",
                    "items": {
                        "type": "object",
                        "required": ["row", "text"],
                        "properties": {
                            "row": {"type": "integer", "description": "0-based screen row"},
                            "text": {"type": "string"}
                        }
                    }
                },
                "cursor": {
                    "type": "object",
                    "description": "Cursor position, 0-based (screen = changes)",
                    "properties": {"row": {"type": "integer"}, "col": {"type": "integer"}}
                },
                "matched": {"type": "boolean", "description": "Whether wait_for matched (only when wait_for was given)"},
                "session": {"type": "string", "description": "Session this result is for (multi-session tools)"},
                "sessions": {
                    "type": "array",
                    "description": "Every open session and its activity (multi-session tools)",
                    "items": {"type": "object", "properties": {"name": {"type": "string"}, "activity": {"type": "string"}}}
                },
                "restarted": {"type": "boolean", "description": "The session had exited and a fresh shell was started for this call"},
                "closed": {"type": "boolean", "description": "The session was closed"},
                "activity": {
                    "type": "string",
                    "enum": ["prompt", "continuation", "running", "input", "exited"],
                    "description": "What the session is doing. input: a running command (read, select, a y/n question) waits for a typed line"
                },
                "input_prompt": {"type": "string", "description": "The question on the cursor line while activity = input, such as 'Continue? [y/N] '"},
                "running_command": {"type": "string", "description": "Command line still running (activity = running)"},
                "exit_code": {"type": "integer", "description": "Shell exit code (activity = exited)"},
                "waiting_for_input": {"type": "boolean", "description": "False when wait_ms ran out while a command was still working"},
                "full_screen": {"type": "boolean", "description": "A full-screen program such as vi or less is open"},
                "commands": {
                    "type": "array",
                    "description": "Commands that finished during this call",
                    "items": {
                        "type": "object",
                        "required": ["command", "output", "exit_code"],
                        "properties": {
                            "command": {"type": "string"},
                            "output": {"type": "string"},
                            "exit_code": {"type": "integer"},
                            "output_truncated": {"type": "boolean"}
                        }
                    }
                }
            }
        })
    }

    /// OpenAI-compatible function definition.
    pub fn tool_definition(&self) -> Value {
        json!({
            "type": "function",
            "function": {
                "name": self.name(),
                "description": self.description(),
                "parameters": self.input_schema(),
            }
        })
    }

    /// Send the keys in `args.input`, run until the session needs input (or
    /// `wait_ms` passes), and report the screen and finished commands.
    pub async fn call(&mut self, args: Value) -> Result<Value, TerminalToolError> {
        let Some(args) = args.as_object() else {
            return Err(TerminalToolError("arguments must be a JSON object".into()));
        };
        let input = match args.get("input") {
            None | Some(Value::Null) => "",
            Some(Value::String(s)) => s.as_str(),
            Some(_) => return Err(TerminalToolError("`input` must be a string".into())),
        };
        if input.len() > MAX_INPUT_BYTES {
            return Err(TerminalToolError(format!(
                "`input` is longer than {MAX_INPUT_BYTES} bytes; send it in smaller parts"
            )));
        }
        let wait_ms = match args.get("wait_ms") {
            None | Some(Value::Null) => DEFAULT_WAIT_MS,
            Some(v) => v.as_u64().ok_or_else(|| {
                TerminalToolError("`wait_ms` must be a non-negative integer".into())
            })?,
        }
        .min(MAX_WAIT_MS);
        let wait_for = match args.get("wait_for") {
            None | Some(Value::Null) => None,
            Some(Value::String(p)) => Some(compile_pattern(p)?),
            Some(_) => return Err(TerminalToolError("`wait_for` must be a string".into())),
        };
        let mode = match args.get("screen") {
            None | Some(Value::Null) => ScreenMode::Full,
            Some(Value::String(m)) if m == "full" => ScreenMode::Full,
            Some(Value::String(m)) if m == "changes" => ScreenMode::Changes,
            Some(Value::String(m)) if m == "none" => ScreenMode::None,
            Some(_) => {
                return Err(TerminalToolError(
                    "`screen` must be \"full\", \"changes\" or \"none\"".into(),
                ));
            }
        };

        let name = match args.get("session") {
            None | Some(Value::Null) => MAIN_SESSION,
            Some(Value::String(s)) => s.as_str(),
            Some(_) => return Err(TerminalToolError("`session` must be a string".into())),
        };
        let close = match args.get("close") {
            None | Some(Value::Null) => false,
            Some(Value::Bool(b)) => *b,
            Some(_) => return Err(TerminalToolError("`close` must be a boolean".into())),
        };
        let idx = self.open_session(name, close)?;
        let Some(idx) = idx else {
            let mut out = json!({"session": name, "closed": true});
            out["sessions"] = self.sessions_json();
            return Ok(out);
        };
        let mut restarted = false;
        if !input.is_empty()
            && matches!(
                self.sessions[idx].terminal.activity(),
                TerminalActivity::Exited(_)
            )
            && let Some(fresh) = self.new_terminal()
        {
            self.sessions[idx].terminal = fresh;
            self.sessions[idx].last_rows = None;
            restarted = true;
        }

        let bytes = parse_keys(input);
        let terminal = &self.sessions[idx].terminal;
        let mark = terminal.tty().output_mark();
        if terminal.send(&bytes) < bytes.len() {
            return Err(TerminalToolError(
                "terminal input buffer is full; wait for the running command first".into(),
            ));
        }
        let wait = Duration::from_millis(wait_ms);
        let (waiting, matched) = match wait_for {
            None => (self.wait_idle(idx, wait).await, None),
            Some(re) => {
                let (waiting, matched) = self.wait_match(idx, wait, mark, &re).await;
                (waiting, Some(matched))
            }
        };
        let mut out = self.report(idx, waiting, mode);
        if let Some(matched) = matched {
            out["matched"] = json!(matched);
        }
        if restarted {
            out["restarted"] = json!(true);
        }
        Ok(out)
    }

    /// Index of session `name`, opening it if needed. `Ok(None)` when it was
    /// closed (`close: true`).
    fn open_session(
        &mut self,
        name: &str,
        close: bool,
    ) -> Result<Option<usize>, TerminalToolError> {
        if name.is_empty()
            || name.len() > MAX_SESSION_NAME
            || !name
                .chars()
                .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
        {
            return Err(TerminalToolError(format!(
                "`session` must be 1-{MAX_SESSION_NAME} letters, digits, '-' or '_'"
            )));
        }
        let found = self.sessions.iter().position(|s| s.name == name);
        if close {
            return match found {
                Some(0) => Err(TerminalToolError(
                    "the main session cannot be closed".into(),
                )),
                Some(i) => {
                    self.sessions.remove(i);
                    Ok(None)
                }
                None => Err(TerminalToolError(format!("no session named '{name}'"))),
            };
        }
        if let Some(i) = found {
            return Ok(Some(i));
        }
        if self.factory.is_none() {
            return Err(TerminalToolError(format!(
                "no session named '{name}'; this terminal has only the main session"
            )));
        }
        if self.sessions.len() >= MAX_SESSIONS {
            return Err(TerminalToolError(format!(
                "at most {MAX_SESSIONS} sessions; close one first"
            )));
        }
        let terminal = self.new_terminal().expect("factory checked above");
        self.sessions.push(ToolSession::new(name, terminal));
        Ok(Some(self.sessions.len() - 1))
    }

    /// Run session `idx` until it needs input or `wait` passes, while the
    /// other sessions keep running in the background.
    async fn wait_idle(&mut self, idx: usize, wait: Duration) -> bool {
        let (before, rest) = self.sessions.split_at_mut(idx);
        let (target, after) = rest.split_first_mut().expect("session index in range");
        let others = futures_util::future::join_all(
            before
                .iter_mut()
                .chain(after.iter_mut())
                .map(|s| s.terminal.run_until_idle()),
        );
        let background = async {
            others.await;
            std::future::pending::<()>().await;
        };
        tokio::select! {
            r = crate::time_compat::timeout(wait, target.terminal.run_until_idle()) => {
                matches!(r, Ok(TerminalStatus::Idle | TerminalStatus::Exited(_)))
            }
            () = background => unreachable!("background never completes"),
        }
    }

    /// Run until `re` matches output written since `mark`, the session needs
    /// input, or `wait` passes. Returns (waiting_for_input, matched).
    async fn wait_match(
        &mut self,
        idx: usize,
        wait: Duration,
        mark: u64,
        re: &regex::Regex,
    ) -> (bool, bool) {
        let deadline = crate::time_compat::Instant::now() + wait;
        loop {
            if self.output_matches(idx, mark, re) {
                return (self.sessions[idx].terminal.tty().is_idle(), true);
            }
            let left = deadline.saturating_duration_since(crate::time_compat::Instant::now());
            if left.is_zero() {
                return (false, false);
            }
            if self.wait_idle(idx, left.min(WAIT_SLICE)).await {
                return (true, self.output_matches(idx, mark, re));
            }
        }
    }

    fn output_matches(&self, idx: usize, mark: u64, re: &regex::Regex) -> bool {
        let raw = self.sessions[idx].terminal.tty().output_since(mark);
        re.is_match(&strip_ansi(&String::from_utf8_lossy(&raw)))
    }

    fn sessions_json(&self) -> Value {
        json!(
            self.sessions
                .iter()
                .map(|s| json!({"name": s.name, "activity": activity_name(&s.terminal)}))
                .collect::<Vec<_>>()
        )
    }

    fn report(&mut self, idx: usize, waiting: bool, mode: ScreenMode) -> Value {
        let multi = self.sessions.len() > 1;
        let sessions = multi.then(|| self.sessions_json());
        let session = &mut self.sessions[idx];
        let terminal = &session.terminal;
        let mut out = json!({
            "waiting_for_input": waiting,
            "full_screen": terminal.is_alternate_screen(),
            "commands": terminal
                .take_transcript()
                .into_iter()
                .map(record_json)
                .collect::<Vec<_>>(),
        });
        let (activity, extra) = match terminal.activity() {
            TerminalActivity::Starting | TerminalActivity::Prompt => ("prompt", None),
            TerminalActivity::ContinuationPrompt => ("continuation", None),
            TerminalActivity::Running { command } => {
                if let Some(prompt) = terminal.input_prompt() {
                    out["input_prompt"] = json!(prompt);
                    ("input", Some(("running_command", json!(command))))
                } else {
                    ("running", Some(("running_command", json!(command))))
                }
            }
            TerminalActivity::Exited(code) => ("exited", Some(("exit_code", json!(code)))),
        };
        out["activity"] = json!(activity);
        if let Some((key, value)) = extra {
            out[key] = value;
        }
        let rows = terminal.screen_rows();
        match mode {
            ScreenMode::Full => out["screen"] = json!(terminal.screen_text()),
            ScreenMode::Changes => {
                let prev = session.last_rows.as_deref().unwrap_or(&[]);
                let resized = prev.len() != rows.len();
                let changes: Vec<Value> = rows
                    .iter()
                    .enumerate()
                    .filter(|(i, row)| resized || prev.get(*i) != Some(*row))
                    .map(|(i, row)| json!({"row": i, "text": row}))
                    .collect();
                let (row, col) = terminal.cursor();
                out["screen_changes"] = json!(changes);
                out["cursor"] = json!({"row": row, "col": col});
            }
            ScreenMode::None => {}
        }
        session.last_rows = Some(rows);
        if let Some(sessions) = sessions {
            out["session"] = json!(session.name);
            out["sessions"] = sessions;
        }
        out
    }
}

/// The `activity` value reported for a terminal.
fn activity_name(terminal: &Terminal) -> &'static str {
    match terminal.activity() {
        TerminalActivity::Starting | TerminalActivity::Prompt => "prompt",
        TerminalActivity::ContinuationPrompt => "continuation",
        TerminalActivity::Running { .. } if terminal.input_prompt().is_some() => "input",
        TerminalActivity::Running { .. } => "running",
        TerminalActivity::Exited(_) => "exited",
    }
}

fn compile_pattern(pattern: &str) -> Result<regex::Regex, TerminalToolError> {
    if pattern.len() > MAX_PATTERN_BYTES {
        return Err(TerminalToolError(format!(
            "`wait_for` is longer than {MAX_PATTERN_BYTES} bytes"
        )));
    }
    regex::RegexBuilder::new(pattern)
        .size_limit(MAX_REGEX_SIZE)
        .multi_line(true)
        .build()
        .map_err(|e| TerminalToolError(format!("`wait_for` is not a valid regex: {e}")))
}

/// Drop terminal control sequences (CSI, OSC, two-byte ESC) and `\r` so
/// `wait_for` matches the text a reader sees.
fn strip_ansi(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut chars = s.chars().peekable();
    while let Some(c) = chars.next() {
        match c {
            '\r' => {}
            '\x1b' => match chars.next() {
                Some('[') => {
                    for c in chars.by_ref() {
                        if ('@'..='~').contains(&c) {
                            break;
                        }
                    }
                }
                Some(']') => {
                    while let Some(c) = chars.next() {
                        if c == '\x07' {
                            break;
                        }
                        if c == '\x1b' && chars.peek() == Some(&'\\') {
                            chars.next();
                            break;
                        }
                    }
                }
                _ => {}
            },
            c => out.push(c),
        }
    }
    out
}

fn record_json(r: CommandRecord) -> Value {
    let mut v = json!({
        "command": r.command,
        "output": r.output,
        "exit_code": r.exit_code,
    });
    if r.output_truncated {
        v["output_truncated"] = json!(true);
    }
    v
}

/// Translate Vim key notation into terminal input bytes. Unknown `<...>`
/// tokens are typed literally.
pub(crate) fn parse_keys(input: &str) -> Vec<u8> {
    let mut out = Vec::with_capacity(input.len());
    let mut rest = input;
    while let Some(start) = rest.find('<') {
        out.extend_from_slice(&rest.as_bytes()[..start]);
        let after = &rest[start + 1..];
        let token = after
            .find('>')
            .map(|end| &after[..end])
            .filter(|t| !t.is_empty() && t.len() <= 10 && !t.contains('<'));
        match token.and_then(key_bytes) {
            Some(bytes) => {
                out.extend_from_slice(&bytes);
                rest = &after[token.map_or(0, str::len) + 1..];
            }
            None => {
                out.push(b'<');
                rest = after;
            }
        }
    }
    out.extend_from_slice(rest.as_bytes());
    out
}

fn key_bytes(token: &str) -> Option<Vec<u8>> {
    let lower = token.to_ascii_lowercase();
    let fixed: &[u8] = match lower.as_str() {
        "enter" | "cr" | "return" => b"\r",
        "esc" | "escape" => b"\x1b",
        "tab" => b"\t",
        "bs" | "backspace" => b"\x7f",
        "del" | "delete" => b"\x1b[3~",
        "space" => b" ",
        "lt" => b"<",
        "up" => b"\x1b[A",
        "down" => b"\x1b[B",
        "right" => b"\x1b[C",
        "left" => b"\x1b[D",
        "home" => b"\x1b[H",
        "end" => b"\x1b[F",
        "pageup" => b"\x1b[5~",
        "pagedown" => b"\x1b[6~",
        _ => {
            let ctrl = lower.strip_prefix("c-")?;
            let [c] = ctrl.as_bytes() else {
                return None;
            };
            return match c {
                b'a'..=b'z' => Some(vec![c - b'a' + 1]),
                b'[' => Some(vec![0x1b]),
                _ => None,
            };
        }
    };
    Some(fixed.to_vec())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Bash;

    #[test]
    fn vim_notation_maps_to_bytes() {
        assert_eq!(parse_keys("ls<Enter>"), b"ls\r");
        assert_eq!(parse_keys("ihi<Esc>:wq<CR>"), b"ihi\x1b:wq\r");
        assert_eq!(parse_keys("<C-c><c-D>"), b"\x03\x04");
        assert_eq!(parse_keys("<Up><PageDown>"), b"\x1b[A\x1b[6~");
        assert_eq!(parse_keys("a <lt> b"), b"a < b");
    }

    #[test]
    fn unknown_or_unclosed_tokens_are_literal() {
        assert_eq!(parse_keys("echo <foo>"), b"echo <foo>");
        assert_eq!(parse_keys("a<b"), b"a<b");
        assert_eq!(parse_keys("x << EOF<Enter>"), b"x << EOF\r");
        assert_eq!(parse_keys("<C-1>"), b"<C-1>");
        assert_eq!(parse_keys("<>"), b"<>");
    }

    #[tokio::test]
    async fn edits_a_file_in_vi_over_several_calls() {
        let mut tool = TerminalTool::new(Bash::builder());
        let out = tool
            .call(json!({"input": "vi /tmp/n.txt<Enter>"}))
            .await
            .unwrap();
        assert_eq!(out["activity"], "running");
        assert_eq!(out["running_command"], "vi /tmp/n.txt");
        assert_eq!(out["full_screen"], true);
        assert_eq!(out["waiting_for_input"], true);

        let out = tool
            .call(json!({"input": "ihello<Esc>:wq<Enter>"}))
            .await
            .unwrap();
        assert_eq!(out["activity"], "prompt");
        assert_eq!(out["full_screen"], false);
        assert_eq!(out["commands"][0]["command"], "vi /tmp/n.txt");
        assert_eq!(out["commands"][0]["exit_code"], 0);
        let saved = tool
            .terminal()
            .fs()
            .read_file("/tmp/n.txt".as_ref())
            .await
            .unwrap();
        assert_eq!(saved, b"hello\n");
    }

    #[tokio::test]
    async fn slow_command_reports_running_then_finishes() {
        let mut tool = TerminalTool::new(Bash::builder());
        let out = tool
            .call(json!({"input": "sleep 0.2; echo done<Enter>", "wait_ms": 20}))
            .await
            .unwrap();
        assert_eq!(out["waiting_for_input"], false);
        assert_eq!(out["activity"], "running");
        let out = tool.call(json!({"wait_ms": 2000})).await.unwrap();
        assert_eq!(out["waiting_for_input"], true);
        assert_eq!(out["commands"][0]["output"], "done\n");
    }

    #[tokio::test]
    async fn exit_is_reported() {
        let mut tool = TerminalTool::new(Bash::builder());
        let out = tool.call(json!({"input": "exit 5<Enter>"})).await.unwrap();
        assert_eq!(out["activity"], "exited");
        assert_eq!(out["exit_code"], 5);
    }

    #[tokio::test]
    async fn bad_arguments_are_user_facing_errors() {
        let mut tool = TerminalTool::new(Bash::builder());
        assert!(tool.call(json!("ls")).await.is_err());
        assert!(tool.call(json!({"input": 3})).await.is_err());
        assert!(
            tool.call(json!({"input": "", "wait_ms": -1}))
                .await
                .is_err()
        );
        let big = "a".repeat(MAX_INPUT_BYTES + 1);
        let err = tool.call(json!({"input": big})).await.unwrap_err();
        assert!(err.to_string().contains("smaller parts"));
    }

    #[tokio::test]
    async fn wait_for_returns_on_match_while_running() {
        let mut tool = TerminalTool::new(Bash::builder());
        let started = std::time::Instant::now(); // std-time-ok: native-only test
        let out = tool
            .call(json!({
                "input": "for i in 1 2 3; do echo tick$i; sleep 1; done<Enter>",
                "wait_for": "tick[2]",
                "wait_ms": 10000
            }))
            .await
            .unwrap();
        assert_eq!(out["matched"], true);
        assert_eq!(out["activity"], "running");
        assert_eq!(out["waiting_for_input"], false);
        assert!(started.elapsed() < std::time::Duration::from_millis(2500));
        assert!(out["screen"].as_str().unwrap().contains("tick2"));
        let out = tool.call(json!({"input": "<C-c>"})).await.unwrap();
        assert_eq!(out["activity"], "prompt");
    }

    #[tokio::test]
    async fn wait_for_ignores_typed_echo() {
        let mut tool = TerminalTool::new(Bash::builder());
        let out = tool
            .call(json!({"input": "echo ready<Enter>", "wait_for": "^ready$"}))
            .await
            .unwrap();
        assert_eq!(out["matched"], true);
        let out = tool
            .call(json!({"input": "true ready<Enter>", "wait_for": "ready"}))
            .await
            .unwrap();
        assert_eq!(out["matched"], false);
        assert_eq!(out["waiting_for_input"], true);
    }

    #[tokio::test]
    async fn wait_for_times_out_without_match() {
        let mut tool = TerminalTool::new(Bash::builder());
        let out = tool
            .call(json!({"input": "sleep 5<Enter>", "wait_for": "never", "wait_ms": 100}))
            .await
            .unwrap();
        assert_eq!(out["matched"], false);
        assert_eq!(out["activity"], "running");
        assert!(out.get("matched").is_some());
        let out = tool.call(json!({"input": "<C-c>"})).await.unwrap();
        assert!(out.get("matched").is_none());
    }

    #[test]
    fn wait_for_rejects_bad_patterns() {
        assert!(
            compile_pattern("(")
                .unwrap_err()
                .to_string()
                .contains("not a valid regex")
        );
        assert!(compile_pattern(&"a".repeat(MAX_PATTERN_BYTES + 1)).is_err());
        assert!(compile_pattern("a{1000}{1000}").is_err());
    }

    #[test]
    fn strip_ansi_keeps_visible_text() {
        assert_eq!(strip_ansi("\x1b[1;31mred\x1b[0m ok\r\n"), "red ok\n");
        assert_eq!(strip_ansi("\x1b]0;title\x07x\x1b]2;t\x1b\\y\x1b7z"), "xyz");
    }

    #[tokio::test]
    async fn screen_changes_returns_only_changed_rows() {
        let mut tool = TerminalTool::new(Bash::builder());
        let out = tool
            .call(json!({"input": "echo one<Enter>", "screen": "changes"}))
            .await
            .unwrap();
        // First call: every row is new.
        assert_eq!(out["screen_changes"].as_array().unwrap().len(), 24);
        assert!(out.get("screen").is_none());
        let out = tool
            .call(json!({"input": "echo two<Enter>", "screen": "changes"}))
            .await
            .unwrap();
        let changes = out["screen_changes"].as_array().unwrap();
        let texts: Vec<_> = changes
            .iter()
            .map(|c| c["text"].as_str().unwrap())
            .collect();
        assert_eq!(texts, ["$ echo two", "two", "$"]);
        assert_eq!(changes[0]["row"], 2);
        assert_eq!(out["cursor"], json!({"row": 4, "col": 2}));
        let out = tool.call(json!({"screen": "none"})).await.unwrap();
        assert!(out.get("screen").is_none() && out.get("screen_changes").is_none());
        let out = tool.call(json!({"screen": "changes"})).await.unwrap();
        assert_eq!(out["screen_changes"], json!([]));
        assert!(tool.call(json!({"screen": "diff"})).await.is_err());
        assert!(tool.call(json!({"wait_for": 1})).await.is_err());
    }

    #[tokio::test]
    async fn question_is_reported_as_input() {
        let mut tool = TerminalTool::new(Bash::builder());
        let out = tool
            .call(json!({"input": "read -p 'Overwrite? [y/N] ' a; echo a=$a<Enter>"}))
            .await
            .unwrap();
        assert_eq!(out["activity"], "input");
        assert_eq!(out["input_prompt"], "Overwrite? [y/N] ");
        assert_eq!(out["waiting_for_input"], true);
        let out = tool.call(json!({"input": "y<Enter>"})).await.unwrap();
        assert_eq!(out["activity"], "prompt");
        assert!(out.get("input_prompt").is_none());
        assert_eq!(out["commands"][0]["output"], "a=y\n");
        // vi is running, not asking a line question.
        let out = tool.call(json!({"input": "vi<Enter>"})).await.unwrap();
        assert_eq!(out["activity"], "running");
    }

    fn multi() -> TerminalTool {
        TerminalTool::with_sessions(TerminalSize::default(), Box::new(Bash::builder))
    }

    #[tokio::test]
    async fn sessions_share_files_not_shell_state() {
        let mut tool = multi();
        tool.call(json!({"input": "cd /tmp; X=1; echo data > f<Enter>"}))
            .await
            .unwrap();
        let out = tool
            .call(json!({"session": "b", "input": "cat /tmp/f; echo x=$X; pwd<Enter>"}))
            .await
            .unwrap();
        assert_eq!(out["session"], "b");
        assert_eq!(out["commands"][0]["output"], "data\nx=\n/home/user\n");
        let names: Vec<_> = out["sessions"]
            .as_array()
            .unwrap()
            .iter()
            .map(|s| s["name"].as_str().unwrap())
            .collect();
        assert_eq!(names, ["main", "b"]);
        assert_eq!(tool.session_names(), ["main", "b"]);
        // Back on main, its state is intact.
        let out = tool
            .call(json!({"input": "echo x=$X; pwd<Enter>"}))
            .await
            .unwrap();
        assert_eq!(out["commands"][0]["output"], "x=1\n/tmp\n");
    }

    #[tokio::test]
    async fn other_sessions_keep_running() {
        let mut tool = multi();
        // A slow writer in "bg"; return right away.
        let out = tool
            .call(json!({
                "session": "bg",
                "input": "for i in 1 2 3; do echo $i >> /tmp/log; sleep 0.1; done<Enter>",
                "wait_ms": 0
            }))
            .await
            .unwrap();
        assert_eq!(out["activity"], "running");
        // Waiting in main drives bg too.
        let out = tool
            .call(json!({"input": "sleep 0.6; cat /tmp/log<Enter>"}))
            .await
            .unwrap();
        assert_eq!(out["commands"][0]["output"], "1\n2\n3\n");
        let bg = &out["sessions"][1];
        assert_eq!(bg["name"], "bg");
        assert_eq!(bg["activity"], "prompt");
    }

    #[tokio::test]
    async fn close_restart_and_limits() {
        let mut tool = multi();
        tool.call(json!({"session": "t", "input": "exit 3<Enter>"}))
            .await
            .unwrap();
        let out = tool
            .call(json!({"session": "t", "input": "echo back<Enter>"}))
            .await
            .unwrap();
        assert_eq!(out["restarted"], true);
        assert_eq!(out["commands"][0]["output"], "back\n");
        let out = tool
            .call(json!({"session": "t", "close": true}))
            .await
            .unwrap();
        assert_eq!(out["closed"], true);
        assert_eq!(tool.session_names(), ["main"]);
        assert!(tool.call(json!({"close": true})).await.is_err());
        assert!(tool.call(json!({"session": "a b"})).await.is_err());
        assert!(
            tool.call(json!({"session": "nope", "close": true}))
                .await
                .is_err()
        );
        for i in 1..MAX_SESSIONS {
            tool.call(json!({"session": format!("s{i}"), "wait_ms": 0}))
                .await
                .unwrap();
        }
        let err = tool.call(json!({"session": "extra"})).await.unwrap_err();
        assert!(err.to_string().contains("at most 8"));
    }

    #[tokio::test]
    async fn single_session_tool_rejects_other_sessions() {
        let mut tool = TerminalTool::new(Bash::builder());
        assert!(tool.input_schema()["properties"].get("session").is_none());
        let err = tool.call(json!({"session": "b"})).await.unwrap_err();
        assert!(err.to_string().contains("only the main session"));
        let out = tool
            .call(json!({"session": "main", "input": "true<Enter>"}))
            .await
            .unwrap();
        assert!(out.get("sessions").is_none());
        assert!(
            multi().input_schema()["properties"]
                .get("session")
                .is_some()
        );
        assert!(multi().system_prompt().contains("session"));
    }

    #[test]
    fn definition_matches_schema() {
        let tool = TerminalTool::new(Bash::builder());
        let def = tool.tool_definition();
        assert_eq!(def["function"]["name"], "terminal");
        assert_eq!(def["function"]["parameters"], tool.input_schema());
        assert!(tool.system_prompt().starts_with("terminal:"));
    }
}