DS4Server/src/agent.rs

use crate::model::ModelChoice;
use rfd::{MessageButtons, MessageDialog, MessageDialogResult, MessageLevel};
use serde_json::{Map, Value};
use std::collections::{HashMap, HashSet};
use std::ffi::{CStr, CString, c_char, c_int, c_void};
use std::fs::{self, File};
use std::os::unix::process::CommandExt;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

const MAX_FILE_BYTES: u64 = 16 * 1024 * 1024;

#[repr(C)]
struct WebConfig {
    home_dir: *const c_char,
    port: c_int,
    confirm: Option<unsafe extern "C" fn(*mut c_void, *const c_char, *mut c_char, usize) -> c_int>,
    confirm_data: *mut c_void,
    log: Option<unsafe extern "C" fn(*mut c_void, *const c_char)>,
    log_data: *mut c_void,
    cancel: Option<unsafe extern "C" fn(*mut c_void) -> bool>,
    cancel_data: *mut c_void,
}

unsafe extern "C" {
    fn ds4_web_create(config: *const WebConfig) -> *mut c_void;
    fn ds4_web_free(web: *mut c_void);
    fn ds4_web_google_search(
        web: *mut c_void,
        query: *const c_char,
        error: *mut c_char,
        error_len: usize,
    ) -> *mut c_char;
    fn ds4_web_visit_page(
        web: *mut c_void,
        url: *const c_char,
        error: *mut c_char,
        error_len: usize,
    ) -> *mut c_char;
    fn free(pointer: *mut c_void);
}

struct WebCallbacks {
    cancel: AtomicPtr<AtomicBool>,
}

struct Browser {
    web: *mut c_void,
    callbacks: Box<WebCallbacks>,
}

// The C browser is used by one agent-tools worker at a time under Tools' mutex.
unsafe impl Send for Browser {}

impl Browser {
    fn new() -> Result<Self, String> {
        let mut callbacks = Box::new(WebCallbacks {
            cancel: AtomicPtr::new(std::ptr::null_mut()),
        });
        let home = CString::new(
            std::env::var_os("HOME")
                .unwrap_or_else(|| ".".into())
                .to_string_lossy()
                .as_bytes(),
        )
        .map_err(|_| "The home directory contains a NUL byte.".to_owned())?;
        let data = (&mut *callbacks) as *mut WebCallbacks as *mut c_void;
        let config = WebConfig {
            home_dir: home.as_ptr(),
            port: 9333,
            confirm: Some(web_confirm),
            confirm_data: data,
            log: None,
            log_data: std::ptr::null_mut(),
            cancel: Some(web_cancel),
            cancel_data: data,
        };
        let web = unsafe { ds4_web_create(&config) };
        if web.is_null() {
            Err("Could not initialize browser tools.".into())
        } else {
            Ok(Self { web, callbacks })
        }
    }

    fn google_search(&self, query: &str, cancel: &AtomicBool) -> Result<String, String> {
        self.call(query, cancel, ds4_web_google_search)
    }

    fn visit_page(&self, url: &str, cancel: &AtomicBool) -> Result<String, String> {
        self.call(url, cancel, ds4_web_visit_page)
    }

    fn call(
        &self,
        value: &str,
        cancel: &AtomicBool,
        operation: unsafe extern "C" fn(
            *mut c_void,
            *const c_char,
            *mut c_char,
            usize,
        ) -> *mut c_char,
    ) -> Result<String, String> {
        let value = CString::new(value).map_err(|_| "web input contains a NUL byte".to_owned())?;
        let mut error = [0 as c_char; 256];
        self.callbacks.cancel.store(
            cancel as *const AtomicBool as *mut AtomicBool,
            Ordering::Relaxed,
        );
        let result =
            unsafe { operation(self.web, value.as_ptr(), error.as_mut_ptr(), error.len()) };
        self.callbacks
            .cancel
            .store(std::ptr::null_mut(), Ordering::Relaxed);
        if result.is_null() {
            let error = unsafe { CStr::from_ptr(error.as_ptr()) }
                .to_string_lossy()
                .into_owned();
            Err(if error.is_empty() {
                "browser tool failed".into()
            } else {
                error
            })
        } else {
            let output = unsafe { CStr::from_ptr(result) }
                .to_string_lossy()
                .into_owned();
            unsafe { free(result.cast()) };
            Ok(output)
        }
    }
}

impl Drop for Browser {
    fn drop(&mut self) {
        unsafe { ds4_web_free(self.web) };
    }
}

unsafe extern "C" fn web_confirm(
    _data: *mut c_void,
    message: *const c_char,
    error: *mut c_char,
    error_len: usize,
) -> c_int {
    let message = if message.is_null() {
        "The web tool wants to start a visible Chrome browser. Allow?".into()
    } else {
        unsafe { CStr::from_ptr(message) }.to_string_lossy()
    };
    let allowed = MessageDialog::new()
        .set_level(MessageLevel::Warning)
        .set_title("Allow browser tools?")
        .set_description(message)
        .set_buttons(MessageButtons::YesNo)
        .show()
        == MessageDialogResult::Yes;
    if !allowed {
        write_c_error(error, error_len, "user denied Chrome browser start");
    }
    c_int::from(allowed)
}

unsafe extern "C" fn web_cancel(data: *mut c_void) -> bool {
    if data.is_null() {
        return false;
    }
    let callbacks = unsafe { &*(data.cast::<WebCallbacks>()) };
    let cancel = callbacks.cancel.load(Ordering::Relaxed);
    !cancel.is_null() && unsafe { (*cancel).load(Ordering::Relaxed) }
}

fn write_c_error(output: *mut c_char, length: usize, message: &str) {
    if output.is_null() || length == 0 {
        return;
    }
    let bytes = message.as_bytes();
    let count = bytes.len().min(length - 1);
    unsafe {
        std::ptr::copy_nonoverlapping(bytes.as_ptr().cast(), output, count);
        *output.add(count) = 0;
    }
}

const TOOL_SCHEMAS: &str = r#"{"type":"function","function":{"name":"google_search","description":"Search Google in a visible browser and return compact Markdown links.","parameters":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}}
{"type":"function","function":{"name":"visit_page","description":"Open a URL in a visible browser and return rendered page text.","parameters":{"type":"object","properties":{"url":{"type":"string"}},"required":["url"]}}}
{"type":"function","function":{"name":"bash","description":"Run a shell command.","parameters":{"type":"object","properties":{"command":{"type":"string"},"timeout_sec":{"type":"number"},"refresh_sec":{"type":"number"}},"required":["command"]}}}
{"type":"function","function":{"name":"bash_status","description":"Report current status and new output for a bash job.","parameters":{"type":"object","properties":{"job":{"type":"number"},"pid":{"type":"number"},"refresh_sec":{"type":"number"}},"required":["job"]}}}
{"type":"function","function":{"name":"bash_stop","description":"Terminate a running bash job and report its final output.","parameters":{"type":"object","properties":{"job":{"type":"number"},"pid":{"type":"number"},"refresh_sec":{"type":"number"}},"required":["job"]}}}
{"type":"function","function":{"name":"read","description":"Read a text file or a range of lines.","parameters":{"type":"object","properties":{"path":{"type":"string"},"start_line":{"type":"number"},"max_lines":{"type":"number"},"whole":{"type":"boolean"},"raw":{"type":"boolean"}},"required":["path"]}}}
{"type":"function","function":{"name":"more","description":"Continue the previous read-like output.","parameters":{"type":"object","properties":{"count":{"type":"number"}}}}}
{"type":"function","function":{"name":"write","description":"Create or overwrite a text file.","parameters":{"type":"object","properties":{"path":{"type":"string"},"content":{"type":"string"}},"required":["path","content"]}}}
{"type":"function","function":{"name":"edit","description":"Replace exactly one old text match; old may contain [upto] between unique head and tail anchors.","parameters":{"type":"object","properties":{"path":{"type":"string"},"old":{"type":"string"},"new":{"type":"string"}},"required":["path","old","new"]}}}
{"type":"function","function":{"name":"search","description":"Search files and return compact edit-friendly matches.","parameters":{"type":"object","properties":{"query":{"type":"string"},"path":{"type":"string"},"mode":{"type":"string"},"glob":{"type":"string"},"context":{"type":"number"},"max_results":{"type":"number"},"case_sensitive":{"type":"boolean"}},"required":["query"]}}}
{"type":"function","function":{"name":"list","description":"List one directory compactly.","parameters":{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}}}"#;

#[derive(Clone, Debug, PartialEq)]
pub(crate) struct ToolCall {
    pub(crate) name: String,
    pub(crate) arguments: Map<String, Value>,
}

pub(crate) struct ActiveTools {
    pub(crate) results: Receiver<String>,
    pub(crate) cancel: Arc<AtomicBool>,
}

struct BashJob {
    id: u32,
    command: String,
    child: Child,
    output: PathBuf,
    started: Instant,
    timeout: Duration,
    observed: usize,
}

pub(crate) struct Tools {
    root: PathBuf,
    context_tokens: i32,
    more: Option<(PathBuf, usize, bool)>,
    jobs: HashMap<u32, BashJob>,
    temporary_files: HashSet<PathBuf>,
    next_job: u32,
    browser: Browser,
}

impl Tools {
    pub(crate) fn new(root: &Path, context_tokens: i32) -> Result<Self, String> {
        Ok(Self {
            root: root
                .canonicalize()
                .map_err(|error| format!("Could not open the project directory: {error}"))?,
            context_tokens,
            more: None,
            jobs: HashMap::new(),
            temporary_files: HashSet::new(),
            next_job: 1,
            browser: Browser::new()?,
        })
    }

    fn execute(&mut self, call: &ToolCall, cancel: &AtomicBool) -> String {
        let result = match call.name.as_str() {
            "read" => self.read(call),
            "more" => self.more(call),
            "write" => self.write(call),
            "edit" => self.edit(call),
            "search" => self.search(call),
            "list" => self.list(call),
            "bash" => self.bash(call, cancel),
            "bash_status" => self.bash_observe(call, false, cancel),
            "bash_stop" => self.bash_observe(call, true, cancel),
            "google_search" => self.google_search(call, cancel),
            "visit_page" => self.visit_page(call, cancel),
            name => Err(format!("unknown tool: {name}")),
        };
        match result {
            Ok(result) if result.len() <= self.result_limit() => result,
            Ok(result) => format!(
                "Tool error: {} result is too large for this context ({} bytes). Retry with a smaller read/search/bash output.\n",
                call.name,
                result.len()
            ),
            Err(error) => format!("Tool error: {error}\n"),
        }
    }

    fn result_limit(&self) -> usize {
        (self.context_tokens.max(4096) as usize * 2).min(512 * 1024)
    }

    fn existing_path(&self, value: &str) -> Result<PathBuf, String> {
        let path = if Path::new(value).is_absolute() {
            PathBuf::from(value)
        } else {
            self.root.join(value)
        };
        let path = path
            .canonicalize()
            .map_err(|error| format!("open {value}: {error}"))?;
        if self.temporary_files.contains(&path) {
            Ok(path)
        } else {
            self.inside_project(path, value)
        }
    }

    fn writable_path(&self, value: &str) -> Result<PathBuf, String> {
        let path = if Path::new(value).is_absolute() {
            PathBuf::from(value)
        } else {
            self.root.join(value)
        };
        if path.exists() {
            return self.existing_path(value);
        }
        let parent = path
            .parent()
            .ok_or_else(|| format!("invalid path: {value}"))?
            .canonicalize()
            .map_err(|error| format!("open parent of {value}: {error}"))?;
        self.inside_project(parent, value)?;
        Ok(path)
    }

    fn inside_project(&self, path: PathBuf, original: &str) -> Result<PathBuf, String> {
        if path.starts_with(&self.root) {
            Ok(path)
        } else {
            Err(format!("path is outside the project: {original}"))
        }
    }

    fn default_lines(&self) -> usize {
        match self.context_tokens {
            ..=8192 => 120,
            8193..=16384 => 240,
            _ => 500,
        }
    }

    fn read(&mut self, call: &ToolCall) -> Result<String, String> {
        let path = required_string(call, "path")?;
        let start = integer(call, "start_line", 1, 1, usize::MAX);
        let count = integer(call, "max_lines", self.default_lines(), 1, usize::MAX);
        self.read_range(
            &self.existing_path(path)?,
            start,
            count,
            boolean(call, "whole", false),
            boolean(call, "raw", false),
        )
    }

    fn more(&mut self, call: &ToolCall) -> Result<String, String> {
        let (path, start, raw) = self
            .more
            .clone()
            .ok_or_else(|| "no previous output to continue".to_owned())?;
        let count = integer(call, "count", self.default_lines(), 1, usize::MAX);
        self.read_range(&path, start, count, false, raw)
    }

    fn read_range(
        &mut self,
        path: &Path,
        start: usize,
        count: usize,
        whole: bool,
        raw: bool,
    ) -> Result<String, String> {
        let metadata = path.metadata().map_err(|error| error.to_string())?;
        if metadata.len() > MAX_FILE_BYTES {
            return Err(format!(
                "file too large: {} exceeds {MAX_FILE_BYTES} bytes",
                path.display()
            ));
        }
        let data = fs::read_to_string(path)
            .map_err(|error| format!("read {}: {error}", path.display()))?;
        let lines = data.lines().collect::<Vec<_>>();
        let first = start.saturating_sub(1).min(lines.len());
        let last = if whole {
            lines.len()
        } else {
            first.saturating_add(count).min(lines.len())
        };
        self.more = (last < lines.len()).then(|| (path.to_owned(), last + 1, raw));
        let mut output = String::new();
        if !raw {
            if last < lines.len() {
                output.push_str(&format!(
                    "{}: lines {}-{} of {}; continue_offset={}; call more with count={} to read the next chunk\n",
                    path.display(),
                    if lines.is_empty() { 0 } else { first + 1 },
                    last,
                    lines.len(),
                    last + 1,
                    count
                ));
            } else {
                output.push_str(&format!(
                    "{}: lines {}-{} of {}\n",
                    path.display(),
                    if lines.is_empty() { 0 } else { first + 1 },
                    last,
                    lines.len()
                ));
            }
        }
        for (index, line) in lines[first..last].iter().enumerate() {
            if raw {
                output.push_str(line);
            } else {
                output.push_str(&format!("{} {line}", first + index + 1));
            }
            output.push('\n');
        }
        if raw && last < lines.len() {
            output.push_str(&format!(
                "[Read truncated at line {} of {}. continue_offset={}. Call more with count={} to read the next chunk.]\n",
                last,
                lines.len(),
                last + 1,
                count
            ));
        }
        Ok(output)
    }

    fn write(&self, call: &ToolCall) -> Result<String, String> {
        let display = required_string(call, "path")?;
        let content = required_string(call, "content")?;
        if content.len() as u64 > MAX_FILE_BYTES {
            return Err(format!("content exceeds {MAX_FILE_BYTES} bytes"));
        }
        let path = self.writable_path(display)?;
        fs::write(&path, content).map_err(|error| format!("write {display}: {error}"))?;
        Ok(format!("Wrote {} bytes to {display}\n", content.len()))
    }

    fn edit(&self, call: &ToolCall) -> Result<String, String> {
        let display = required_string(call, "path")?;
        let old = required_string(call, "old")?;
        let new = required_string(call, "new")?;
        if old.is_empty() {
            return Err("edit requires non-empty old text".into());
        }
        let path = self.existing_path(display)?;
        if path.metadata().map_err(|error| error.to_string())?.len() > MAX_FILE_BYTES {
            return Err(format!(
                "file too large: {display} exceeds {MAX_FILE_BYTES} bytes"
            ));
        }
        let data = fs::read_to_string(&path).map_err(|error| format!("read {display}: {error}"))?;
        let (start, end, anchored) = edit_span(&data, old)?;
        if data.len() - (end - start) + new.len() > MAX_FILE_BYTES as usize {
            return Err(format!("edited file would exceed {MAX_FILE_BYTES} bytes"));
        }
        let mut output = String::with_capacity(data.len() - (end - start) + new.len());
        output.push_str(&data[..start]);
        output.push_str(new);
        output.push_str(&data[end..]);
        fs::write(&path, output).map_err(|error| format!("write {display}: {error}"))?;
        Ok(format!(
            "Edited {display} using {} replacement\n",
            if anchored {
                "anchored old/new"
            } else {
                "old/new"
            }
        ))
    }

    fn list(&self, call: &ToolCall) -> Result<String, String> {
        let display = string(call, "path").unwrap_or(".");
        let path = self.existing_path(display)?;
        if !path.is_dir() {
            return Err(format!("not a directory: {display}"));
        }
        let mut entries = fs::read_dir(&path)
            .map_err(|error| format!("list {display}: {error}"))?
            .filter_map(Result::ok)
            .collect::<Vec<_>>();
        entries.sort_by_key(|entry| entry.file_name());
        let mut output = format!("{display}:\n");
        for entry in entries.iter().take(300) {
            let metadata = fs::symlink_metadata(entry.path()).map_err(|error| error.to_string())?;
            let kind = if metadata.file_type().is_symlink() {
                'l'
            } else if metadata.is_dir() {
                'd'
            } else {
                '-'
            };
            let suffix = if metadata.is_dir() { "/" } else { "" };
            output.push_str(&format!(
                "{kind} {:>10} {}{suffix}\n",
                metadata.len(),
                entry.file_name().to_string_lossy()
            ));
        }
        if entries.len() > 300 {
            output.push_str("... more entries omitted ...\n");
        }
        Ok(output)
    }

    fn search(&self, call: &ToolCall) -> Result<String, String> {
        let query = required_string(call, "query")?;
        let display = string(call, "path").unwrap_or(".");
        let path = self.existing_path(display)?;
        let context = integer(call, "context", 0, 0, 5);
        let limit = integer(call, "max_results", 50, 1, 500);
        let options = SearchOptions {
            query,
            glob: string(call, "glob"),
            regex: string(call, "mode") == Some("regex"),
            case_sensitive: boolean(call, "case_sensitive", true),
            context,
            limit,
        };
        search_path(&self.root, &path, &options)
    }

    fn bash(&mut self, call: &ToolCall, cancel: &AtomicBool) -> Result<String, String> {
        let command = required_string(call, "command")?.to_owned();
        let timeout = integer(call, "timeout_sec", 3600, 1, 24 * 3600);
        let refresh = integer(call, "refresh_sec", 60, 1, 3600);
        let id = self.next_job;
        self.next_job = self.next_job.saturating_add(1).max(1);
        let output = std::env::temp_dir().join(format!(
            "ds4_agent_output_{}_{}_{}",
            std::process::id(),
            id,
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos()
        ));
        let stdout = File::create(&output).map_err(|error| error.to_string())?;
        let output = output.canonicalize().map_err(|error| error.to_string())?;
        self.temporary_files.insert(output.clone());
        let stderr = stdout.try_clone().map_err(|error| error.to_string())?;
        let child = Command::new("/bin/sh")
            .arg("-c")
            .arg(&command)
            .current_dir(&self.root)
            .stdin(Stdio::null())
            .stdout(stdout)
            .stderr(stderr)
            .process_group(0)
            .spawn()
            .map_err(|error| format!("bash failed to start: {error}"))?;
        self.jobs.insert(
            id,
            BashJob {
                id,
                command,
                child,
                output,
                started: Instant::now(),
                timeout: Duration::from_secs(timeout as u64),
                observed: 0,
            },
        );
        self.wait_job(id, refresh, cancel, true)
    }

    fn bash_observe(
        &mut self,
        call: &ToolCall,
        stop: bool,
        cancel: &AtomicBool,
    ) -> Result<String, String> {
        let requested_id = integer(call, "job", 0, 0, u32::MAX as usize) as u32;
        let requested_pid = integer(call, "pid", 0, 0, u32::MAX as usize) as u32;
        let id = if self.jobs.contains_key(&requested_id) {
            requested_id
        } else {
            self.jobs
                .iter()
                .find_map(|(id, job)| (job.child.id() == requested_pid).then_some(*id))
                .unwrap_or(requested_id)
        };
        let refresh = integer(call, "refresh_sec", 60, 1, 3600);
        if !self.jobs.contains_key(&id) {
            return Err(format!(
                "bash job not found: job={requested_id} pid={requested_pid}"
            ));
        }
        if stop {
            stop_job(self.jobs.get_mut(&id).unwrap());
        }
        self.wait_job(id, if stop { 1 } else { refresh }, cancel, stop)
    }

    fn wait_job(
        &mut self,
        id: u32,
        refresh: usize,
        cancel: &AtomicBool,
        remove_done: bool,
    ) -> Result<String, String> {
        let deadline = Instant::now() + Duration::from_secs(refresh as u64);
        loop {
            let job = self.jobs.get_mut(&id).unwrap();
            let done = job.child.try_wait().map_err(|error| error.to_string())?;
            if done.is_some() || Instant::now() >= deadline {
                break;
            }
            if job.started.elapsed() >= job.timeout {
                stop_job(job);
                break;
            }
            if cancel.load(Ordering::Relaxed) {
                stop_job(job);
                return Err("interrupted".into());
            }
            thread::sleep(Duration::from_millis(100));
        }
        let job = self.jobs.get_mut(&id).unwrap();
        let status = job.child.try_wait().map_err(|error| error.to_string())?;
        let bytes = fs::read(&job.output).unwrap_or_default();
        let first_observation = job.observed == 0;
        let new = &bytes[job.observed.min(bytes.len())..];
        job.observed = bytes.len();
        let shown = if first_observation && new.len() > 8 * 1024 {
            &new[..8 * 1024]
        } else if new.len() > 32 * 1024 {
            &new[new.len() - 32 * 1024..]
        } else {
            new
        };
        let mut result = format!(
            "bash job={} pid={} status={} command={}\noutput_path={}\n",
            job.id,
            job.child.id(),
            status.map_or("running".into(), |status| status.to_string()),
            job.command,
            job.output.display()
        );
        result.push_str(&String::from_utf8_lossy(shown));
        if shown.len() < new.len() {
            if status.is_some() {
                let tail = &new[new.len().saturating_sub(32 * 1024).max(shown.len())..];
                result.push_str("\n... middle output omitted ...\n");
                result.push_str(&String::from_utf8_lossy(tail));
            } else {
                result.push_str("\n... output truncated; use bash_status for new output ...\n");
            }
        }
        if !result.ends_with('\n') {
            result.push('\n');
        }
        if remove_done && status.is_some() {
            let job = self.jobs.remove(&id).unwrap();
            self.temporary_files.remove(&job.output);
            let _ = fs::remove_file(job.output);
        }
        Ok(result)
    }

    fn google_search(&mut self, call: &ToolCall, cancel: &AtomicBool) -> Result<String, String> {
        let query = required_string(call, "query")?;
        self.browser.google_search(query, cancel)
    }

    fn visit_page(&mut self, call: &ToolCall, cancel: &AtomicBool) -> Result<String, String> {
        let url = required_string(call, "url")?;
        if !matches!(
            url.split_once(':').map(|part| part.0),
            Some("http" | "https")
        ) {
            return Err("visit_page requires an HTTP or HTTPS URL".into());
        }
        let markdown = self.browser.visit_page(url, cancel)?;
        let output = std::env::temp_dir().join(format!(
            "ds4_agent_web_{}_{}",
            std::process::id(),
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos()
        ));
        fs::write(&output, &markdown).map_err(|error| error.to_string())?;
        let output = output.canonicalize().map_err(|error| error.to_string())?;
        self.temporary_files.insert(output.clone());
        let head = markdown
            .lines()
            .take(100)
            .collect::<Vec<_>>()
            .join("\n")
            .chars()
            .take(8 * 1024)
            .collect::<String>();
        Ok(format!(
            "visit_page url={url}\noutput_path={} ({} bytes, {} lines)\n<head -100 {}>\n{head}\n</head>\nUse read with raw=true to inspect more rendered text.\n",
            output.display(),
            markdown.len(),
            markdown.lines().count(),
            output.display()
        ))
    }
}

struct SearchOptions<'a> {
    query: &'a str,
    glob: Option<&'a str>,
    regex: bool,
    case_sensitive: bool,
    context: usize,
    limit: usize,
}

fn search_path(root: &Path, path: &Path, options: &SearchOptions<'_>) -> Result<String, String> {
    let mut files = Vec::new();
    collect_search_files(path, 0, &mut files)?;
    let mut matches = 0;
    let mut body = String::new();
    for file in files {
        if matches >= options.limit {
            break;
        }
        let relative = file.strip_prefix(root).unwrap_or(&file);
        if options.glob.is_some_and(|glob| {
            !wildcard_match(glob, &relative.to_string_lossy())
                && !wildcard_match(
                    glob,
                    &file.file_name().unwrap_or_default().to_string_lossy(),
                )
        }) {
            continue;
        }
        let remaining = options.limit - matches;
        let (count, text) = if options.regex {
            regex_search_file(&file, options, remaining)?
        } else {
            literal_search_file(&file, options, remaining)?
        };
        if count > 0 {
            body.push_str(&format!("{}\n{text}\n", relative.display()));
            matches += count;
        }
    }
    if matches == 0 {
        Ok("No matches\n".into())
    } else {
        Ok(format!(
            "{matches} match{} shown\n\n{body}",
            if matches == 1 { "" } else { "es" }
        ))
    }
}

fn collect_search_files(
    path: &Path,
    depth: usize,
    output: &mut Vec<PathBuf>,
) -> Result<(), String> {
    if depth > 24 {
        return Ok(());
    }
    let metadata = fs::symlink_metadata(path).map_err(|error| error.to_string())?;
    if metadata.file_type().is_symlink() {
        return Ok(());
    }
    if metadata.is_file() {
        if metadata.len() <= MAX_FILE_BYTES {
            output.push(path.to_owned());
        }
        return Ok(());
    }
    if !metadata.is_dir() {
        return Ok(());
    }
    let mut entries = fs::read_dir(path)
        .map_err(|error| error.to_string())?
        .filter_map(Result::ok)
        .collect::<Vec<_>>();
    entries.sort_by_key(|entry| entry.file_name());
    for entry in entries {
        if entry.file_name() == ".git" {
            continue;
        }
        collect_search_files(&entry.path(), depth + 1, output)?;
    }
    Ok(())
}

fn literal_search_file(
    path: &Path,
    options: &SearchOptions<'_>,
    limit: usize,
) -> Result<(usize, String), String> {
    let bytes = fs::read(path).map_err(|error| error.to_string())?;
    if bytes.contains(&0) {
        return Ok((0, String::new()));
    }
    let data = String::from_utf8_lossy(&bytes);
    let lines = data.lines().collect::<Vec<_>>();
    let query = (!options.case_sensitive).then(|| options.query.to_lowercase());
    let found = lines
        .iter()
        .enumerate()
        .filter_map(|(index, line)| {
            let matched = query.as_ref().map_or_else(
                || line.contains(options.query),
                |query| line.to_lowercase().contains(query),
            );
            matched.then_some(index)
        })
        .take(limit)
        .collect::<Vec<_>>();
    let mut output = String::new();
    let mut last = None;
    for match_index in &found {
        let start = match_index.saturating_sub(options.context);
        let end = match_index
            .saturating_add(options.context + 1)
            .min(lines.len());
        for (index, line) in lines.iter().enumerate().take(end).skip(start) {
            if last.is_none_or(|last| index > last) {
                output.push_str(&format!("  {} {line}\n", index + 1));
                last = Some(index);
            }
        }
    }
    Ok((found.len(), output))
}

fn regex_search_file(
    path: &Path,
    options: &SearchOptions<'_>,
    limit: usize,
) -> Result<(usize, String), String> {
    let mut command = Command::new("/usr/bin/grep");
    command.arg("-n").arg("-I").arg("-E");
    if !options.case_sensitive {
        command.arg("-i");
    }
    if options.context > 0 {
        command.arg("-C").arg(options.context.to_string());
    }
    command.arg("-m").arg(limit.to_string());
    command.arg("--").arg(options.query).arg(path);
    let output = command.output().map_err(|error| error.to_string())?;
    if output.status.code() == Some(1) {
        return Ok((0, String::new()));
    }
    if !output.status.success() {
        return Err(format!(
            "invalid regex: {}",
            String::from_utf8_lossy(&output.stderr).trim()
        ));
    }
    let output = String::from_utf8_lossy(&output.stdout).into_owned();
    let count = output
        .lines()
        .filter(|line| {
            line.split_once(':')
                .is_some_and(|(number, _)| number.parse::<usize>().is_ok())
        })
        .count();
    Ok((count, output))
}

fn wildcard_match(pattern: &str, value: &str) -> bool {
    let (pattern, value) = (pattern.as_bytes(), value.as_bytes());
    let (mut p, mut v, mut star, mut retry) = (0, 0, None, 0);
    while v < value.len() {
        if p < pattern.len() && (pattern[p] == b'?' || pattern[p] == value[v]) {
            p += 1;
            v += 1;
        } else if p < pattern.len() && pattern[p] == b'*' {
            star = Some(p);
            p += 1;
            retry = v;
        } else if let Some(star) = star {
            p = star + 1;
            retry += 1;
            v = retry;
        } else {
            return false;
        }
    }
    while p < pattern.len() && pattern[p] == b'*' {
        p += 1;
    }
    p == pattern.len()
}

impl Drop for Tools {
    fn drop(&mut self) {
        for job in self.jobs.values_mut() {
            stop_job(job);
            let _ = fs::remove_file(&job.output);
        }
        for path in &self.temporary_files {
            let _ = fs::remove_file(path);
        }
    }
}

pub(crate) fn execute_async(tools: Arc<Mutex<Tools>>, calls: Vec<ToolCall>) -> ActiveTools {
    let cancel = Arc::new(AtomicBool::new(false));
    let worker_cancel = Arc::clone(&cancel);
    let (sender, results) = mpsc::channel();
    thread::Builder::new()
        .name("agent-tools".into())
        .spawn(move || {
            let mut output = String::new();
            let mut tools = tools
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            for (index, call) in calls.iter().enumerate() {
                if worker_cancel.load(Ordering::Relaxed) {
                    output.push_str("Tool error: interrupted\n");
                    break;
                }
                output.push_str(&format!("Tool result {} ({}):\n", index + 1, call.name));
                output.push_str(&tools.execute(call, &worker_cancel));
                if !output.ends_with('\n') {
                    output.push('\n');
                }
            }
            let _ = sender.send(output);
        })
        .expect("agent tool worker must start");
    ActiveTools { results, cancel }
}

pub(crate) fn error_async(error: String) -> ActiveTools {
    let cancel = Arc::new(AtomicBool::new(false));
    let (sender, results) = mpsc::channel();
    let _ = sender.send(format!(
        "Tool error: invalid tool call: {error}\nRetry using the exact tool syntax from the system prompt.\n"
    ));
    ActiveTools { results, cancel }
}

pub(crate) fn parse_tool_calls(
    model: ModelChoice,
    text: &str,
) -> Result<(String, Vec<ToolCall>), String> {
    let (content, calls) = if model == ModelChoice::Glm52 {
        parse_glm_calls(text)?
    } else {
        crate::server::parse_dsml_tool_calls(text)?
    };
    calls
        .into_iter()
        .map(|(name, arguments)| match arguments {
            Value::Object(arguments) => Ok(ToolCall { name, arguments }),
            _ => Err(format!("tool {name} arguments are not an object")),
        })
        .collect::<Result<Vec<_>, _>>()
        .map(|calls| (content, calls))
}

pub(crate) fn system_prompt(model: ModelChoice, extra: &str) -> String {
    let tools = if model == ModelChoice::Glm52 {
        format!(
            "You are a coding agent running in a local workspace. Use tools for local file and system work. Avoid printing large file contents or large code blocks as answers; create or edit files with tools, then summarize results briefly.\n\n# Tools\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{TOOL_SCHEMAS}\n</tools>\n\nFor a function call, output exactly: <tool_call>function-name<arg_key>key</arg_key><arg_value>value</arg_value></tool_call>\nTool calls are not allowed inside <think></think>. Use read/search for focused context, edit with exact unique old text, and [upto] only between unique head and tail anchors. Use refresh_sec for long bash jobs and poll with bash_status or stop with bash_stop. Preserve the current system configuration unless the user explicitly asks otherwise."
        )
    } else {
        format!(
            "You are a coding agent running in a local workspace. Use tools for local file and system work. Avoid printing large file contents or large code blocks as answers; create or edit files with tools, then summarize results briefly.\n\n## Tools\n\nInvoke native DSML tools exactly as:\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"$TOOL_NAME\">\n<｜DSML｜parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE</｜DSML｜parameter>\n</｜DSML｜invoke>\n</｜DSML｜tool_calls>\n\nTool calls are not allowed inside <think></think>. String parameters use raw text and string=\"true\"; numbers and booleans use JSON text and string=\"false\". Read defaults to a bounded chunk; use more to continue and whole=true only when needed. Use write for new files or whole-file replacement. Use edit with path first and exact unique old text; old may contain one [upto] marker between unique head and tail anchors. For long bash commands pass refresh_sec, then use bash_status or bash_stop. The first web call asks permission to start visible Chrome.\n\n### Available Tool Schemas\n\n{TOOL_SCHEMAS}\n\n# Rules\n- Always use strict DSML syntax.\n- Use read/search to get anchors before editing.\n- Preserve the current system configuration unless explicitly asked otherwise."
        )
    };
    if extra.trim().is_empty() {
        tools
    } else {
        format!("{tools}\n\n{extra}")
    }
}

pub(crate) fn try_tool_result(active: &ActiveTools) -> Result<Option<String>, String> {
    match active.results.try_recv() {
        Ok(result) => Ok(Some(result)),
        Err(TryRecvError::Empty) => Ok(None),
        Err(TryRecvError::Disconnected) => {
            Err("The agent tool worker stopped unexpectedly.".into())
        }
    }
}

pub(crate) fn visible_content(content: &str) -> &str {
    [
        "<｜DSML｜tool_calls>",
        "<DSML｜tool_calls>",
        "<tool_calls>",
        "<tool_call>",
    ]
    .into_iter()
    .filter_map(|marker| content.find(marker))
    .min()
    .map_or(content, |end| content[..end].trim_end())
}

pub(crate) fn tool_summaries(model: ModelChoice, content: &str) -> Vec<String> {
    parse_tool_calls(model, content)
        .map(|(_, calls)| {
            calls
                .into_iter()
                .map(|call| {
                    let detail = ["path", "command", "query", "url"]
                        .into_iter()
                        .find_map(|name| string(&call, name))
                        .map(|value| {
                            let mut value = value.replace('\n', " ");
                            if value.chars().count() > 120 {
                                value = value.chars().take(119).collect::<String>() + "…";
                            }
                            format!("  {value}")
                        })
                        .unwrap_or_default();
                    format!("🛠 {}{detail}", call.name)
                })
                .collect()
        })
        .unwrap_or_default()
}

fn parse_glm_calls(text: &str) -> Result<(String, Vec<(String, Value)>), String> {
    let scan = text
        .rfind("</think>")
        .map_or(text, |position| &text[position + "</think>".len()..]);
    let Some(first) = scan.find("<tool_call>") else {
        return Ok((text.to_owned(), Vec::new()));
    };
    let visible_len = text.len() - scan.len() + first;
    let mut rest = &scan[first..];
    let mut calls = Vec::new();
    while rest.starts_with("<tool_call>") {
        rest = &rest["<tool_call>".len()..];
        let end = rest
            .find("</tool_call>")
            .ok_or_else(|| "incomplete GLM tool call".to_owned())?;
        let body = &rest[..end];
        let name_end = body.find("<arg_key>").unwrap_or(body.len());
        let name = body[..name_end].trim();
        if name.is_empty() {
            return Err("GLM tool call without function name".into());
        }
        let mut arguments = Map::new();
        let mut args = &body[name_end..];
        while !args.is_empty() {
            let key = between(&mut args, "<arg_key>", "</arg_key>")?;
            let value = between(&mut args, "<arg_value>", "</arg_value>")?;
            arguments.insert(key.to_owned(), Value::String(value.to_owned()));
        }
        calls.push((name.to_owned(), Value::Object(arguments)));
        rest = rest[end + "</tool_call>".len()..].trim_start();
    }
    Ok((text[..visible_len].trim_end().to_owned(), calls))
}

fn between<'a>(input: &mut &'a str, open: &str, close: &str) -> Result<&'a str, String> {
    let body = input
        .strip_prefix(open)
        .ok_or_else(|| format!("expected {open}"))?;
    let end = body
        .find(close)
        .ok_or_else(|| format!("expected {close}"))?;
    *input = &body[end + close.len()..];
    Ok(&body[..end])
}

fn edit_span(data: &str, old: &str) -> Result<(usize, usize, bool), String> {
    let markers = old.match_indices("[upto]").collect::<Vec<_>>();
    if markers.len() > 1 {
        return Err("old text contains more than one [upto] marker".into());
    }
    if let Some((marker, _)) = markers.first() {
        let head = &old[..*marker];
        let tail = old[*marker + "[upto]".len()..].trim_start_matches(['\r', '\n']);
        if tail.trim().is_empty() {
            return Err("old text after [upto] must include a unique tail anchor".into());
        }
        let start = unique_match(data, head, "old head")?;
        let after_head = start + head.len();
        let tail_offset = unique_match(&data[after_head..], tail, "old tail")?;
        return Ok((start, after_head + tail_offset + tail.len(), true));
    }
    let start = unique_match(data, old, "old text")?;
    Ok((start, start + old.len(), false))
}

fn unique_match(data: &str, needle: &str, label: &str) -> Result<usize, String> {
    if needle.is_empty() {
        return Err(format!("{label} is empty"));
    }
    let matches = data
        .match_indices(needle)
        .map(|(index, _)| index)
        .collect::<Vec<_>>();
    match matches.as_slice() {
        [] => Err(format!("{label} was not found")),
        [index] => Ok(*index),
        _ => Err(format!("{label} is ambiguous ({} matches)", matches.len())),
    }
}

fn stop_job(job: &mut BashJob) {
    if job.child.try_wait().ok().flatten().is_some() {
        return;
    }
    let pid = job.child.id();
    let _ = Command::new("/bin/kill")
        .arg("-TERM")
        .arg(format!("-{pid}"))
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status();
    let deadline = Instant::now() + Duration::from_secs(1);
    while Instant::now() < deadline {
        if job.child.try_wait().ok().flatten().is_some() {
            return;
        }
        thread::sleep(Duration::from_millis(20));
    }
    let _ = Command::new("/bin/kill")
        .arg("-KILL")
        .arg(format!("-{pid}"))
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status();
    let _ = job.child.wait();
}

fn string<'a>(call: &'a ToolCall, name: &str) -> Option<&'a str> {
    call.arguments.get(name).and_then(Value::as_str)
}

fn required_string<'a>(call: &'a ToolCall, name: &str) -> Result<&'a str, String> {
    string(call, name).ok_or_else(|| format!("{} requires {name}", call.name))
}

fn integer(call: &ToolCall, name: &str, default: usize, min: usize, max: usize) -> usize {
    call.arguments
        .get(name)
        .and_then(|value| {
            value
                .as_u64()
                .or_else(|| value.as_i64().map(|value| value.max(0) as u64))
                .or_else(|| {
                    value
                        .as_f64()
                        .filter(|value| value.is_finite())
                        .map(|value| value.max(0.0) as u64)
                })
                .or_else(|| {
                    value
                        .as_str()
                        .and_then(|value| value.parse::<f64>().ok())
                        .filter(|value| value.is_finite())
                        .map(|value| value.max(0.0) as u64)
                })
        })
        .and_then(|value| usize::try_from(value).ok())
        .unwrap_or(default)
        .clamp(min, max)
}

fn boolean(call: &ToolCall, name: &str, default: bool) -> bool {
    call.arguments
        .get(name)
        .and_then(|value| {
            value.as_bool().or_else(|| {
                value.as_str().and_then(|value| {
                    if value.eq_ignore_ascii_case("true")
                        || value.eq_ignore_ascii_case("yes")
                        || value == "1"
                    {
                        Some(true)
                    } else if value.eq_ignore_ascii_case("false")
                        || value.eq_ignore_ascii_case("no")
                        || value == "0"
                    {
                        Some(false)
                    } else {
                        None
                    }
                })
            })
        })
        .unwrap_or(default)
}

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

    #[test]
    fn prompts_and_parsers_expose_the_reference_tool_set() {
        let prompt = system_prompt(ModelChoice::DeepSeekV4Flash, "extra");
        for name in [
            "google_search",
            "visit_page",
            "bash",
            "bash_status",
            "bash_stop",
            "read",
            "more",
            "write",
            "edit",
            "search",
            "list",
        ] {
            assert!(prompt.contains(&format!("\"name\":\"{name}\"")));
        }
        assert!(prompt.ends_with("extra"));

        let text = "done<tool_call>read<arg_key>path</arg_key><arg_value>src/main.rs</arg_value></tool_call>";
        let (visible, calls) = parse_tool_calls(ModelChoice::Glm52, text).unwrap();
        assert_eq!(visible, "done");
        assert_eq!(calls[0].name, "read");
        assert_eq!(calls[0].arguments["path"], "src/main.rs");

        let dsml = "done<｜DSML｜tool_calls><｜DSML｜invoke name=\"read\"><｜DSML｜parameter name=\"path\" string=\"true\">src/main.rs</｜DSML｜parameter></｜DSML｜invoke></｜DSML｜tool_calls>";
        let (visible, calls) = parse_tool_calls(ModelChoice::DeepSeekV4Flash, dsml).unwrap();
        assert_eq!(visible, "done");
        assert_eq!(calls[0].arguments["path"], "src/main.rs");
        assert!(
            parse_tool_calls(
                ModelChoice::DeepSeekV4Flash,
                "<｜DSML｜tool_calls><｜DSML｜invoke name=\"read\">"
            )
            .is_err()
        );
    }

    #[test]
    fn anchored_edits_require_unique_head_and_tail() {
        let data = "start\nold one\nold two\nfinish\nother\n";
        assert_eq!(
            edit_span(data, "start\n[upto]\nfinish\n").unwrap(),
            (0, 29, true)
        );
        assert!(edit_span("same same", "same").is_err());
        assert!(edit_span(data, "start\n[upto]\n").is_err());
    }

    #[test]
    fn project_boundary_rejects_parent_paths() {
        let directory = std::env::temp_dir().join(format!(
            "ds4-agent-test-{}",
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        fs::create_dir_all(&directory).unwrap();
        let tools = Tools::new(&directory, 32_768).unwrap();
        assert!(tools.writable_path("inside.txt").is_ok());
        assert!(tools.writable_path("../outside.txt").is_err());
        fs::remove_dir_all(directory).unwrap();
    }

    #[test]
    fn local_file_and_bash_tools_execute_in_the_project() {
        let directory = std::env::temp_dir().join(format!(
            "ds4-agent-tools-{}",
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        fs::create_dir_all(&directory).unwrap();
        let mut tools = Tools::new(&directory, 4096).unwrap();
        let cancel = AtomicBool::new(false);

        let write = call("write", [("path", "note.txt"), ("content", "one\ntwo\n")]);
        assert!(tools.execute(&write, &cancel).starts_with("Wrote 8 bytes"));
        let read = call("read", [("path", "note.txt"), ("max_lines", "1")]);
        assert!(tools.execute(&read, &cancel).contains("1 one"));
        assert!(tools.execute(&call("more", []), &cancel).contains("2 two"));
        let edit = call(
            "edit",
            [("path", "note.txt"), ("old", "two"), ("new", "three")],
        );
        assert!(tools.execute(&edit, &cancel).starts_with("Edited note.txt"));
        let search = call("search", [("query", "three"), ("glob", "*.txt")]);
        assert!(tools.execute(&search, &cancel).contains("2 three"));
        assert!(
            tools
                .execute(&call("list", [("path", ".")]), &cancel)
                .contains("note.txt")
        );
        let bash = call(
            "bash",
            [("command", "printf shell-ok"), ("refresh_sec", "1")],
        );
        assert!(tools.execute(&bash, &cancel).contains("shell-ok"));
        let running = call(
            "bash",
            [
                ("command", "printf started; sleep 5; printf finished"),
                ("refresh_sec", "1"),
            ],
        );
        assert!(tools.execute(&running, &cancel).contains("status=running"));
        let stopped = tools.execute(&call("bash_stop", [("job", "2")]), &cancel);
        assert!(!stopped.contains("status=running"));
        assert_eq!(
            fs::read_to_string(directory.join("note.txt")).unwrap(),
            "one\nthree\n"
        );
        fs::remove_dir_all(directory).unwrap();
    }

    fn call<const N: usize>(name: &str, arguments: [(&str, &str); N]) -> ToolCall {
        ToolCall {
            name: name.to_owned(),
            arguments: arguments
                .into_iter()
                .map(|(name, value)| (name.to_owned(), Value::String(value.to_owned())))
                .collect(),
        }
    }
}