2150 lines
77 KiB
Rust
2150 lines
77 KiB
Rust
use crate::model::ModelChoice;
|
||
use serde_json::{Map, Value};
|
||
use std::collections::HashMap;
|
||
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, RecvTimeoutError, Sender, 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;
|
||
const SHELL_ENV_ALLOWLIST: &[&str] = &[
|
||
"PATH",
|
||
"HOME",
|
||
"USER",
|
||
"LOGNAME",
|
||
"SHELL",
|
||
"TMPDIR",
|
||
"LANG",
|
||
"LC_ALL",
|
||
"TERM",
|
||
"DEVELOPER_DIR",
|
||
"SDKROOT",
|
||
"MACOSX_DEPLOYMENT_TARGET",
|
||
"RUSTUP_TOOLCHAIN",
|
||
];
|
||
pub(crate) const COMPACTION_OBSERVATION_PREFIX: &str = "Bash job update after context compaction.";
|
||
|
||
#[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;
|
||
let _ = error;
|
||
let _ = error_len;
|
||
// Approval is obtained by the shared tool worker before entering C. Keeping
|
||
// this callback lets the reference browser retain its launch guard.
|
||
1
|
||
}
|
||
|
||
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) }
|
||
}
|
||
|
||
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. Search is literal by default; set mode to regex for patterns such as foo|bar.","parameters":{"type":"object","properties":{"query":{"type":"string"},"path":{"type":"string"},"mode":{"type":"string","enum":["literal","regex"]},"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) events: Receiver<ToolEvent>,
|
||
pub(crate) cancel: Arc<AtomicBool>,
|
||
worker: Option<thread::JoinHandle<()>>,
|
||
}
|
||
|
||
impl Drop for ActiveTools {
|
||
fn drop(&mut self) {
|
||
self.cancel.store(true, Ordering::Relaxed);
|
||
if let Some(worker) = self.worker.take() {
|
||
let _ = worker.join();
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||
pub(crate) enum ToolLifecycle {
|
||
Parsing,
|
||
AwaitingApproval,
|
||
Queued,
|
||
Running,
|
||
Completed,
|
||
Failed,
|
||
Stopped,
|
||
}
|
||
|
||
impl ToolLifecycle {
|
||
pub(crate) fn label(self) -> &'static str {
|
||
match self {
|
||
Self::Parsing => "Parsing",
|
||
Self::AwaitingApproval => "Awaiting approval",
|
||
Self::Queued => "Queued",
|
||
Self::Running => "Running",
|
||
Self::Completed => "Completed",
|
||
Self::Failed => "Failed",
|
||
Self::Stopped => "Stopped",
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Debug)]
|
||
pub(crate) struct ToolCard {
|
||
pub(crate) call: ToolCall,
|
||
pub(crate) state: ToolLifecycle,
|
||
pub(crate) result: Option<String>,
|
||
}
|
||
|
||
impl ToolCard {
|
||
pub(crate) fn parsing(call: ToolCall) -> Self {
|
||
Self {
|
||
call,
|
||
state: ToolLifecycle::Parsing,
|
||
result: None,
|
||
}
|
||
}
|
||
|
||
pub(crate) fn streaming() -> Self {
|
||
Self::parsing(ToolCall {
|
||
name: "Tool call".into(),
|
||
arguments: Map::new(),
|
||
})
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Debug)]
|
||
pub(crate) struct ApprovalPrompt {
|
||
pub(crate) title: String,
|
||
pub(crate) detail: String,
|
||
pub(crate) working_directory: PathBuf,
|
||
}
|
||
|
||
pub(crate) enum ToolEvent {
|
||
State {
|
||
index: usize,
|
||
state: ToolLifecycle,
|
||
result: Option<String>,
|
||
},
|
||
Approval {
|
||
index: usize,
|
||
prompt: ApprovalPrompt,
|
||
decision: Sender<bool>,
|
||
},
|
||
}
|
||
|
||
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)>,
|
||
more_text: Option<(String, usize)>,
|
||
jobs: HashMap<u32, BashJob>,
|
||
next_job: u32,
|
||
browser: Browser,
|
||
dev_brain: Option<crate::dev_brain::DevBrain>,
|
||
}
|
||
|
||
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,
|
||
more_text: None,
|
||
jobs: HashMap::new(),
|
||
next_job: 1,
|
||
browser: Browser::new()?,
|
||
dev_brain: None,
|
||
})
|
||
}
|
||
|
||
pub(crate) fn enable_dev_brain(
|
||
&mut self,
|
||
config: &crate::config::DevBrainConfig,
|
||
projects: &[crate::database::Project],
|
||
) -> Result<(), String> {
|
||
self.dev_brain = Some(crate::dev_brain::DevBrain::open(config, projects)?);
|
||
Ok(())
|
||
}
|
||
|
||
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),
|
||
"dev_brain_info" => self.dev_brain_info(),
|
||
"dev_brain_search" => self.dev_brain_search(call),
|
||
"dev_brain_validate" => self.dev_brain_validate(),
|
||
name => Err(format!("unknown tool: {name}")),
|
||
};
|
||
match result {
|
||
Ok(result) if result.len() <= self.result_limit() => result,
|
||
Ok(result) => {
|
||
let length = result.len();
|
||
self.more_text = Some((result, 0));
|
||
let chunk = self.continue_text(self.default_lines()).unwrap_or_default();
|
||
format!(
|
||
"{} result is too large for this context ({length} bytes); showing a bounded chunk. Use more to continue.\n{chunk}",
|
||
call.name
|
||
)
|
||
}
|
||
Err(error) => format!("Tool error: {error}\n"),
|
||
}
|
||
}
|
||
|
||
fn dev_brain_search(&mut self, call: &ToolCall) -> Result<String, String> {
|
||
let query = required_string(call, "query")?;
|
||
let limit = integer(call, "limit", 8, 1, 50);
|
||
let authoritative = boolean(call, "authoritative", true);
|
||
self.dev_brain
|
||
.as_mut()
|
||
.ok_or_else(|| "Dev Brain is disabled for this session.".to_owned())?
|
||
.search(query, limit, authoritative)
|
||
}
|
||
|
||
fn dev_brain_info(&self) -> Result<String, String> {
|
||
Ok(self
|
||
.dev_brain
|
||
.as_ref()
|
||
.ok_or_else(|| "Dev Brain is disabled for this session.".to_owned())?
|
||
.info())
|
||
}
|
||
|
||
fn dev_brain_validate(&mut self) -> Result<String, String> {
|
||
self.dev_brain
|
||
.as_mut()
|
||
.ok_or_else(|| "Dev Brain is disabled for this session.".to_owned())?
|
||
.validate()
|
||
}
|
||
|
||
fn result_limit(&self) -> usize {
|
||
(self.context_tokens.max(4096) as usize * 2).min(512 * 1024)
|
||
}
|
||
|
||
pub(crate) fn compaction_observation(&mut self) -> Option<String> {
|
||
let mut running = self
|
||
.jobs
|
||
.values_mut()
|
||
.filter_map(|job| {
|
||
job.child.try_wait().ok().flatten().is_none().then(|| {
|
||
format!(
|
||
"bash job={} pid={} status=running command={}\noutput_path={}\n",
|
||
job.id,
|
||
job.child.id(),
|
||
job.command,
|
||
job.output.display()
|
||
)
|
||
})
|
||
})
|
||
.collect::<Vec<_>>();
|
||
if running.is_empty() {
|
||
return None;
|
||
}
|
||
running.sort();
|
||
Some(format!(
|
||
"{COMPACTION_OBSERVATION_PREFIX} Running jobs still need explicit bash_status or bash_stop if relevant.\n{}",
|
||
running.concat()
|
||
))
|
||
}
|
||
|
||
fn existing_path(&self, value: &str) -> Result<PathBuf, String> {
|
||
if Path::new(value)
|
||
.components()
|
||
.any(|component| matches!(component, std::path::Component::ParentDir))
|
||
{
|
||
return Err(format!("path contains a parent traversal: {value}"));
|
||
}
|
||
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}"))?;
|
||
self.inside_readable_root(path, value)
|
||
}
|
||
|
||
fn writable_path(&self, value: &str) -> Result<PathBuf, String> {
|
||
if Path::new(value)
|
||
.components()
|
||
.any(|component| matches!(component, std::path::Component::ParentDir))
|
||
{
|
||
return Err(format!("path contains a parent traversal: {value}"));
|
||
}
|
||
let path = if Path::new(value).is_absolute() {
|
||
PathBuf::from(value)
|
||
} else {
|
||
self.root.join(value)
|
||
};
|
||
if fs::symlink_metadata(&path).is_ok() {
|
||
let path = self.existing_path(value)?;
|
||
return self.inside_writable_root(path, value);
|
||
}
|
||
let mut ancestor = path.as_path();
|
||
let mut suffix = Vec::new();
|
||
while fs::symlink_metadata(ancestor).is_err() {
|
||
let name = ancestor
|
||
.file_name()
|
||
.ok_or_else(|| format!("invalid path: {value}"))?;
|
||
suffix.push(name.to_owned());
|
||
ancestor = ancestor
|
||
.parent()
|
||
.ok_or_else(|| format!("invalid path: {value}"))?;
|
||
}
|
||
let mut resolved = ancestor
|
||
.canonicalize()
|
||
.map_err(|error| format!("open ancestor of {value}: {error}"))?;
|
||
self.inside_readable_root(resolved.clone(), value)?;
|
||
for name in suffix.into_iter().rev() {
|
||
resolved.push(name);
|
||
}
|
||
self.inside_writable_root(resolved, value)
|
||
}
|
||
|
||
fn inside_readable_root(&self, path: PathBuf, original: &str) -> Result<PathBuf, String> {
|
||
if path.starts_with(&self.root) {
|
||
return Ok(path);
|
||
}
|
||
if let Some(brain) = &self.dev_brain
|
||
&& let Ok(relative) = path.strip_prefix(brain.folder())
|
||
&& !relative
|
||
.components()
|
||
.any(|component| component.as_os_str().to_string_lossy().starts_with('.'))
|
||
{
|
||
return Ok(path);
|
||
}
|
||
Err(format!(
|
||
"path is outside the project and managed Dev Brain folder: {original}"
|
||
))
|
||
}
|
||
|
||
fn inside_writable_root(&self, path: PathBuf, original: &str) -> Result<PathBuf, String> {
|
||
if path.starts_with(&self.root)
|
||
|| self
|
||
.dev_brain
|
||
.as_ref()
|
||
.is_some_and(|brain| brain.allows_tool_write(&path))
|
||
{
|
||
Ok(path)
|
||
} else {
|
||
Err(format!(
|
||
"path is not a managed project or Dev Brain file: {original}"
|
||
))
|
||
}
|
||
}
|
||
|
||
fn validate_dev_brain_content(&self, path: &Path, content: &str) -> Result<(), String> {
|
||
if let Some(brain) = &self.dev_brain
|
||
&& path.starts_with(brain.folder())
|
||
{
|
||
brain.validate_tool_content(path, content)?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
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> {
|
||
if self.more_text.is_some() {
|
||
let count = integer(call, "count", self.default_lines(), 1, usize::MAX);
|
||
return Ok(self.continue_text(count).unwrap());
|
||
}
|
||
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 continue_text(&mut self, count: usize) -> Option<String> {
|
||
let (text, start) = self.more_text.take()?;
|
||
let byte_limit = count.saturating_mul(80).min(self.result_limit() / 2);
|
||
let rest = &text[start.min(text.len())..];
|
||
let end = rest
|
||
.char_indices()
|
||
.map(|(index, _)| index)
|
||
.find(|index| *index >= byte_limit)
|
||
.unwrap_or(rest.len());
|
||
let output = rest[..end].to_owned();
|
||
if start + end < text.len() {
|
||
self.more_text = Some((text, start + end));
|
||
}
|
||
Some(output)
|
||
}
|
||
|
||
fn read_range(
|
||
&mut self,
|
||
path: &Path,
|
||
start: usize,
|
||
count: usize,
|
||
whole: bool,
|
||
raw: bool,
|
||
) -> Result<String, String> {
|
||
self.more_text = None;
|
||
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)?;
|
||
self.validate_dev_brain_content(&path, content)?;
|
||
if self
|
||
.dev_brain
|
||
.as_ref()
|
||
.is_some_and(|brain| path.starts_with(brain.folder()))
|
||
&& let Some(parent) = path.parent()
|
||
{
|
||
fs::create_dir_all(parent)
|
||
.map_err(|error| format!("create parent for {display}: {error}"))?;
|
||
}
|
||
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)?;
|
||
self.inside_writable_root(path.clone(), 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..]);
|
||
self.validate_dev_brain_content(&path, &output)?;
|
||
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 dev_brain_path = self
|
||
.dev_brain
|
||
.as_ref()
|
||
.is_some_and(|brain| path.starts_with(brain.folder()));
|
||
let mut output = format!("{display}:\n");
|
||
for entry in entries.iter().take(300) {
|
||
if dev_brain_path && entry.file_name().to_string_lossy().starts_with('.') {
|
||
continue;
|
||
}
|
||
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,
|
||
};
|
||
let root = self
|
||
.dev_brain
|
||
.as_ref()
|
||
.filter(|brain| path.starts_with(brain.folder()))
|
||
.map_or(self.root.as_path(), |brain| brain.folder());
|
||
search_path(root, &path, &options, root != self.root.as_path())
|
||
}
|
||
|
||
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())?;
|
||
let stderr = stdout.try_clone().map_err(|error| error.to_string())?;
|
||
let mut process = Command::new("/bin/sh");
|
||
process
|
||
.arg("-c")
|
||
.arg(format!(
|
||
"ulimit -f {}; exec /bin/sh -c \"$1\"",
|
||
MAX_FILE_BYTES / 512
|
||
))
|
||
.arg("ds4-agent")
|
||
.arg(&command)
|
||
.current_dir(&self.root)
|
||
.stdin(Stdio::null())
|
||
.stdout(stdout)
|
||
.stderr(stderr)
|
||
.process_group(0)
|
||
.env_clear();
|
||
for name in SHELL_ENV_ALLOWLIST {
|
||
if let Some(value) = std::env::var_os(name) {
|
||
process.env(name, value);
|
||
}
|
||
}
|
||
process.env("PWD", &self.root);
|
||
let child = process
|
||
.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 output_limit = self.result_limit().saturating_sub(1024).max(1024);
|
||
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)
|
||
.map_err(|error| format!("read {}: {error}", job.output.display()))?;
|
||
let first_observation = job.observed == 0;
|
||
let new = &bytes[job.observed.min(bytes.len())..];
|
||
job.observed = bytes.len();
|
||
let truncated = new.len() > output_limit;
|
||
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()
|
||
);
|
||
if truncated && status.is_some() {
|
||
let half = output_limit / 2;
|
||
result.push_str(&String::from_utf8_lossy(&new[..half]));
|
||
result.push_str("\n... middle output omitted; open output_path to inspect it ...\n");
|
||
result.push_str(&String::from_utf8_lossy(&new[new.len() - half..]));
|
||
} else {
|
||
result.push_str(&String::from_utf8_lossy(
|
||
&new[..new.len().min(output_limit)],
|
||
));
|
||
}
|
||
if truncated && status.is_none() {
|
||
if first_observation {
|
||
result.push_str("\n... output truncated; open output_path or use bash_status for new output ...\n");
|
||
} else {
|
||
result.push_str(
|
||
"\n... output truncated; open output_path for the complete output ...\n",
|
||
);
|
||
}
|
||
}
|
||
if !result.ends_with('\n') {
|
||
result.push('\n');
|
||
}
|
||
if status.is_some() && !truncated {
|
||
let job = self.jobs.remove(&id).unwrap();
|
||
if let Err(error) = fs::remove_file(&job.output) {
|
||
result.push_str(&format!(
|
||
"Tool warning: could not remove {}: {error}\n",
|
||
job.output.display()
|
||
));
|
||
}
|
||
}
|
||
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 line_end: usize = markdown.split_inclusive('\n').take(100).map(str::len).sum();
|
||
let mut head_end = line_end.min(8 * 1024).min(markdown.len());
|
||
while !markdown.is_char_boundary(head_end) {
|
||
head_end -= 1;
|
||
}
|
||
let head = &markdown[..head_end];
|
||
if head_end < markdown.len() {
|
||
self.more = None;
|
||
self.more_text = Some((markdown.clone(), head_end));
|
||
}
|
||
Ok(format!(
|
||
"visit_page url={url} ({} bytes, {} lines)\n<head -100>\n{head}\n</head>\nPage output is bounded; use more to continue.\n",
|
||
markdown.len(),
|
||
markdown.lines().count()
|
||
))
|
||
}
|
||
|
||
fn approval(&self, call: &ToolCall) -> Option<ApprovalPrompt> {
|
||
if matches!(call.name.as_str(), "google_search" | "visit_page") {
|
||
return Some(ApprovalPrompt {
|
||
title: "Allow visible browser?".into(),
|
||
detail: "Start a visible Chrome browser for this tool call.".into(),
|
||
working_directory: self.root.clone(),
|
||
});
|
||
}
|
||
let command = (call.name == "bash").then(|| string(call, "command"))??;
|
||
risky_shell_reason(command, &self.root).map(|reason| ApprovalPrompt {
|
||
title: "Allow shell command?".into(),
|
||
detail: format!("{reason}\n\n{command}"),
|
||
working_directory: self.root.clone(),
|
||
})
|
||
}
|
||
|
||
pub(crate) fn stop_all_jobs(&mut self) -> Vec<String> {
|
||
let mut failures = Vec::new();
|
||
for job in self.jobs.values_mut() {
|
||
stop_job(job);
|
||
if let Err(error) = fs::remove_file(&job.output) {
|
||
failures.push(format!(
|
||
"Could not remove {}: {error}",
|
||
job.output.display()
|
||
));
|
||
}
|
||
}
|
||
self.jobs.clear();
|
||
failures
|
||
}
|
||
}
|
||
|
||
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<'_>,
|
||
skip_hidden: bool,
|
||
) -> Result<String, String> {
|
||
let mut files = Vec::new();
|
||
collect_search_files(path, 0, skip_hidden, &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,
|
||
skip_hidden: bool,
|
||
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"
|
||
|| skip_hidden && entry.file_name().to_string_lossy().starts_with('.')
|
||
{
|
||
continue;
|
||
}
|
||
collect_search_files(&entry.path(), depth + 1, skip_hidden, 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 failure in self.stop_all_jobs() {
|
||
eprintln!("DS4Server tool cleanup: {failure}");
|
||
}
|
||
}
|
||
}
|
||
|
||
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();
|
||
let (event_sender, events) = mpsc::channel();
|
||
let worker = 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 in 0..calls.len() {
|
||
send_state(&event_sender, index, ToolLifecycle::Queued, None);
|
||
}
|
||
for (index, call) in calls.iter().enumerate() {
|
||
if worker_cancel.load(Ordering::Relaxed) {
|
||
output.push_str("Tool error: interrupted\n");
|
||
send_state(&event_sender, index, ToolLifecycle::Stopped, None);
|
||
break;
|
||
}
|
||
if let Some(prompt) = tools.approval(call) {
|
||
match request_approval(&event_sender, index, prompt, &worker_cancel) {
|
||
Ok(()) => {}
|
||
Err(error) => {
|
||
let result = format!("Tool error: {error}\n");
|
||
output.push_str(&format!(
|
||
"Tool result {} ({}):\n{result}",
|
||
index + 1,
|
||
call.name
|
||
));
|
||
send_state(
|
||
&event_sender,
|
||
index,
|
||
if worker_cancel.load(Ordering::Relaxed) {
|
||
ToolLifecycle::Stopped
|
||
} else {
|
||
ToolLifecycle::Failed
|
||
},
|
||
Some(result),
|
||
);
|
||
if worker_cancel.load(Ordering::Relaxed) {
|
||
break;
|
||
}
|
||
continue;
|
||
}
|
||
}
|
||
}
|
||
send_state(&event_sender, index, ToolLifecycle::Running, None);
|
||
output.push_str(&format!("Tool result {} ({}):\n", index + 1, call.name));
|
||
let result = tools.execute(call, &worker_cancel);
|
||
let state = if worker_cancel.load(Ordering::Relaxed)
|
||
|| result.contains("Tool error: interrupted")
|
||
{
|
||
ToolLifecycle::Stopped
|
||
} else if result.starts_with("Tool error:") {
|
||
ToolLifecycle::Failed
|
||
} else {
|
||
ToolLifecycle::Completed
|
||
};
|
||
output.push_str(&result);
|
||
send_state(&event_sender, index, state, Some(result));
|
||
if !output.ends_with('\n') {
|
||
output.push('\n');
|
||
}
|
||
}
|
||
if worker_cancel.load(Ordering::Relaxed) {
|
||
for failure in tools.stop_all_jobs() {
|
||
output.push_str(&format!("Tool warning: {failure}\n"));
|
||
}
|
||
}
|
||
let _ = sender.send(output);
|
||
})
|
||
.expect("agent tool worker must start");
|
||
ActiveTools {
|
||
results,
|
||
events,
|
||
cancel,
|
||
worker: Some(worker),
|
||
}
|
||
}
|
||
|
||
pub(crate) fn error_async(error: String) -> ActiveTools {
|
||
let cancel = Arc::new(AtomicBool::new(false));
|
||
let (sender, results) = mpsc::channel();
|
||
let (_event_sender, events) = 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,
|
||
events,
|
||
cancel,
|
||
worker: None,
|
||
}
|
||
}
|
||
|
||
fn send_state(
|
||
events: &Sender<ToolEvent>,
|
||
index: usize,
|
||
state: ToolLifecycle,
|
||
result: Option<String>,
|
||
) {
|
||
let _ = events.send(ToolEvent::State {
|
||
index,
|
||
state,
|
||
result,
|
||
});
|
||
}
|
||
|
||
fn request_approval(
|
||
events: &Sender<ToolEvent>,
|
||
index: usize,
|
||
prompt: ApprovalPrompt,
|
||
cancel: &AtomicBool,
|
||
) -> Result<(), String> {
|
||
let (decision, response) = mpsc::channel();
|
||
events
|
||
.send(ToolEvent::Approval {
|
||
index,
|
||
prompt,
|
||
decision,
|
||
})
|
||
.map_err(|_| "approval UI is unavailable".to_owned())?;
|
||
loop {
|
||
if cancel.load(Ordering::Relaxed) {
|
||
return Err("interrupted while awaiting approval".into());
|
||
}
|
||
match response.recv_timeout(Duration::from_millis(50)) {
|
||
Ok(true) => return Ok(()),
|
||
Ok(false) => return Err("user denied this action".into()),
|
||
Err(RecvTimeoutError::Timeout) => {}
|
||
Err(RecvTimeoutError::Disconnected) => {
|
||
return Err("approval UI closed without a decision".into());
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
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, dev_brain: bool) -> String {
|
||
let schemas = if dev_brain {
|
||
format!("{TOOL_SCHEMAS}\n{}", crate::dev_brain::TOOL_SCHEMAS)
|
||
} else {
|
||
TOOL_SCHEMAS.to_owned()
|
||
};
|
||
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{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{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."
|
||
)
|
||
};
|
||
let tools = if dev_brain {
|
||
format!("{tools}\n\n{}", crate::dev_brain::PROMPT)
|
||
} else {
|
||
tools
|
||
};
|
||
if extra.trim().is_empty() {
|
||
tools
|
||
} else {
|
||
format!("{tools}\n\n{extra}")
|
||
}
|
||
}
|
||
|
||
pub(crate) fn system_prompt_reminder(model: ModelChoice, dev_brain: bool) -> String {
|
||
format!(
|
||
"[System prompt reminder follows.]\n{}\n[End system prompt reminder.]",
|
||
system_prompt(model, "", dev_brain)
|
||
)
|
||
}
|
||
|
||
pub(crate) fn prompt_reminder_due(used: u32, last: u32) -> bool {
|
||
used.saturating_sub(last) >= 50_000
|
||
}
|
||
|
||
pub(crate) fn datetime_context() -> String {
|
||
let when = Command::new("/bin/date")
|
||
.arg("+%Y-%m-%d %H:%M:%S %Z")
|
||
.output()
|
||
.ok()
|
||
.filter(|output| output.status.success())
|
||
.map(|output| String::from_utf8_lossy(&output.stdout).trim().to_owned())
|
||
.filter(|output| !output.is_empty())
|
||
.unwrap_or_else(|| {
|
||
SystemTime::now()
|
||
.duration_since(UNIX_EPOCH)
|
||
.unwrap_or_default()
|
||
.as_secs()
|
||
.to_string()
|
||
});
|
||
format!(
|
||
"Current local date and time at session start: {when}. Use this only when date or time matters."
|
||
)
|
||
}
|
||
|
||
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 try_tool_event(active: &ActiveTools) -> Option<ToolEvent> {
|
||
active.events.try_recv().ok()
|
||
}
|
||
|
||
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 has_tool_markup(content: &str) -> bool {
|
||
visible_content(content).len() < content.len()
|
||
}
|
||
|
||
pub(crate) fn stored_tool_cards(
|
||
model: ModelChoice,
|
||
assistant: &str,
|
||
result: Option<&str>,
|
||
) -> Vec<ToolCard> {
|
||
let calls = parse_tool_calls(model, assistant)
|
||
.map(|(_, calls)| calls)
|
||
.unwrap_or_default();
|
||
let results = result.map(split_tool_results).unwrap_or_default();
|
||
calls
|
||
.into_iter()
|
||
.enumerate()
|
||
.map(|(index, call)| {
|
||
let result = results.get(index).cloned();
|
||
let state = result.as_deref().map_or(ToolLifecycle::Stopped, |result| {
|
||
if result.starts_with("Tool error:") {
|
||
ToolLifecycle::Failed
|
||
} else {
|
||
ToolLifecycle::Completed
|
||
}
|
||
});
|
||
ToolCard {
|
||
call,
|
||
state,
|
||
result,
|
||
}
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
fn split_tool_results(result: &str) -> Vec<String> {
|
||
let starts = result
|
||
.match_indices("Tool result ")
|
||
.map(|(index, _)| index)
|
||
.collect::<Vec<_>>();
|
||
starts
|
||
.iter()
|
||
.enumerate()
|
||
.map(|(index, start)| {
|
||
let body = &result[*start..starts.get(index + 1).copied().unwrap_or(result.len())];
|
||
body.split_once("\n")
|
||
.map_or(body, |(_, result)| result)
|
||
.trim_end()
|
||
.to_owned()
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
pub(crate) fn bounded_tool_text(text: &str, limit: usize) -> String {
|
||
let mut output = text.replace('\n', " ");
|
||
if output.chars().count() > limit {
|
||
output = output
|
||
.chars()
|
||
.take(limit.saturating_sub(1))
|
||
.collect::<String>()
|
||
+ "…";
|
||
}
|
||
output
|
||
}
|
||
|
||
pub(crate) fn tool_parameters(call: &ToolCall) -> String {
|
||
["path", "command", "query", "url"]
|
||
.into_iter()
|
||
.find_map(|name| string(call, name).map(|value| (name, value)))
|
||
.map_or_else(
|
||
|| bounded_tool_text(&Value::Object(call.arguments.clone()).to_string(), 240),
|
||
|(name, value)| format!("{name}: {}", bounded_tool_text(value, 240)),
|
||
)
|
||
}
|
||
|
||
pub(crate) fn tool_call_text(call: &ToolCall) -> String {
|
||
let mut tool = Map::new();
|
||
tool.insert("name".into(), Value::String(call.name.clone()));
|
||
tool.insert("arguments".into(), Value::Object(call.arguments.clone()));
|
||
serde_json::to_string_pretty(&Value::Object(tool)).unwrap_or_default()
|
||
}
|
||
|
||
pub(crate) fn tool_output_path(result: &str) -> Option<PathBuf> {
|
||
result
|
||
.lines()
|
||
.find_map(|line| line.strip_prefix("output_path="))
|
||
.map(PathBuf::from)
|
||
.filter(|path| path.is_absolute() && path.is_file())
|
||
}
|
||
|
||
fn risky_shell_reason(command: &str, root: &Path) -> Option<&'static str> {
|
||
let words = command
|
||
.split_whitespace()
|
||
.map(|word| {
|
||
word.trim_matches(|character: char| {
|
||
matches!(
|
||
character,
|
||
'\'' | '"' | ';' | '|' | '&' | '(' | ')' | '{' | '}' | '<' | '>'
|
||
)
|
||
})
|
||
.to_ascii_lowercase()
|
||
})
|
||
.collect::<Vec<_>>();
|
||
if words
|
||
.iter()
|
||
.any(|word| matches!(word.as_str(), "sudo" | "doas" | "su"))
|
||
{
|
||
return Some("This command elevates privileges.");
|
||
}
|
||
if words
|
||
.iter()
|
||
.any(|word| matches!(word.as_str(), "open" | "osascript"))
|
||
{
|
||
return Some("This command launches or controls another application.");
|
||
}
|
||
if words.iter().any(|word| {
|
||
matches!(
|
||
word.as_str(),
|
||
"rm" | "rmdir"
|
||
| "unlink"
|
||
| "shred"
|
||
| "truncate"
|
||
| "dd"
|
||
| "mkfs"
|
||
| "diskutil"
|
||
| "kill"
|
||
| "killall"
|
||
| "pkill"
|
||
)
|
||
}) || (words.iter().any(|word| word == "git")
|
||
&& words
|
||
.iter()
|
||
.any(|word| matches!(word.as_str(), "clean" | "reset")))
|
||
{
|
||
return Some("This command can delete data or discard state.");
|
||
}
|
||
if words.iter().any(|word| {
|
||
matches!(
|
||
word.as_str(),
|
||
"curl" | "wget" | "ssh" | "scp" | "sftp" | "ftp" | "nc" | "ncat" | "telnet"
|
||
)
|
||
}) || (words.iter().any(|word| word == "git")
|
||
&& words
|
||
.iter()
|
||
.any(|word| matches!(word.as_str(), "push" | "pull" | "fetch" | "clone")))
|
||
|| (words.iter().any(|word| {
|
||
matches!(
|
||
word.as_str(),
|
||
"cargo" | "npm" | "pnpm" | "yarn" | "pip" | "pip3"
|
||
)
|
||
}) && words
|
||
.iter()
|
||
.any(|word| matches!(word.as_str(), "install" | "publish" | "login")))
|
||
{
|
||
return Some("This command can create a network side effect.");
|
||
}
|
||
let root = root.to_string_lossy();
|
||
if words.iter().any(|word| {
|
||
(word.contains("../")
|
||
|| word == ".."
|
||
|| word.starts_with("~/")
|
||
|| word.contains("$home")
|
||
|| word.starts_with('/'))
|
||
&& !word.starts_with(root.as_ref())
|
||
&& !matches!(
|
||
word.as_str(),
|
||
"/bin/sh" | "/bin/bash" | "/usr/bin/env" | "/usr/bin/make"
|
||
)
|
||
}) || words.iter().any(|word| {
|
||
matches!(
|
||
word.as_str(),
|
||
"brew" | "launchctl" | "defaults" | "mount" | "umount" | "chown"
|
||
)
|
||
}) {
|
||
return Some("This command can access or change state outside the project.");
|
||
}
|
||
None
|
||
}
|
||
|
||
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", false);
|
||
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"));
|
||
assert!(
|
||
system_prompt_reminder(ModelChoice::DeepSeekV4Flash, false)
|
||
.contains("[System prompt reminder follows.]")
|
||
);
|
||
assert!(!prompt.contains("dev_brain_search"));
|
||
let dev_brain_prompt = system_prompt(ModelChoice::DeepSeekV4Flash, "", true);
|
||
for name in ["dev_brain_info", "dev_brain_search", "dev_brain_validate"] {
|
||
assert!(dev_brain_prompt.contains(name));
|
||
}
|
||
assert!(!dev_brain_prompt.contains("dev_brain_publish"));
|
||
assert!(datetime_context().starts_with("Current local date and time at session start:"));
|
||
assert!(!prompt_reminder_due(49_999, 0));
|
||
assert!(prompt_reminder_due(50_000, 0));
|
||
assert!(!prompt_reminder_due(80_000, 50_000));
|
||
|
||
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 outside = directory.with_extension("outside");
|
||
fs::create_dir_all(&outside).unwrap();
|
||
fs::write(outside.join("secret.txt"), "secret").unwrap();
|
||
std::os::unix::fs::symlink(&outside, directory.join("escape")).unwrap();
|
||
let tools = Tools::new(&directory, 32_768).unwrap();
|
||
assert!(tools.writable_path("inside.txt").is_ok());
|
||
assert!(tools.writable_path("missing/inside.txt").is_ok());
|
||
assert!(tools.writable_path("../outside.txt").is_err());
|
||
assert!(tools.writable_path("missing/../../outside.txt").is_err());
|
||
assert!(tools.existing_path("escape/secret.txt").is_err());
|
||
assert!(tools.writable_path("escape/new.txt").is_err());
|
||
fs::remove_dir_all(directory).unwrap();
|
||
fs::remove_dir_all(outside).unwrap();
|
||
}
|
||
|
||
#[test]
|
||
fn standard_file_tools_can_maintain_only_managed_dev_brain_pages() {
|
||
let directory = std::env::temp_dir().join(format!(
|
||
"ds4-agent-brain-{}",
|
||
SystemTime::now()
|
||
.duration_since(UNIX_EPOCH)
|
||
.unwrap()
|
||
.as_nanos()
|
||
));
|
||
let project = directory.join("project");
|
||
let vault = directory.join("vault");
|
||
fs::create_dir_all(&project).unwrap();
|
||
fs::create_dir_all(vault.join(".obsidian")).unwrap();
|
||
fs::write(project.join("source.rs"), "source\n").unwrap();
|
||
let mut tools = Tools::new(&project, 4096).unwrap();
|
||
tools
|
||
.enable_dev_brain(
|
||
&crate::config::DevBrainConfig {
|
||
enabled: true,
|
||
vault_path: Some(vault.to_string_lossy().into_owned()),
|
||
},
|
||
&[crate::database::Project {
|
||
id: 1,
|
||
name: "Fixture".into(),
|
||
path: project.to_string_lossy().into_owned(),
|
||
collapsed: false,
|
||
}],
|
||
)
|
||
.unwrap();
|
||
|
||
assert!(
|
||
tools
|
||
.existing_path(vault.join("schema.md").to_str().unwrap())
|
||
.is_ok()
|
||
);
|
||
assert!(
|
||
tools
|
||
.writable_path(vault.join("concepts/new.md").to_str().unwrap())
|
||
.is_ok()
|
||
);
|
||
let nested = vault.join("subsystems/inference/modes.md");
|
||
let mut arguments = Map::new();
|
||
arguments.insert(
|
||
"path".into(),
|
||
Value::String(nested.to_string_lossy().into_owned()),
|
||
);
|
||
arguments.insert(
|
||
"content".into(),
|
||
Value::String("---\ndev_brain: true\n---\n# Modes\n".into()),
|
||
);
|
||
tools
|
||
.write(&ToolCall {
|
||
name: "write".into(),
|
||
arguments,
|
||
})
|
||
.unwrap();
|
||
assert!(nested.is_file());
|
||
assert!(
|
||
tools
|
||
.existing_path(vault.join(".obsidian").to_str().unwrap())
|
||
.is_err()
|
||
);
|
||
fs::write(vault.join("concepts/private.md"), "# Private\n").unwrap();
|
||
assert!(
|
||
tools
|
||
.writable_path(vault.join("concepts/private.md").to_str().unwrap())
|
||
.is_err()
|
||
);
|
||
fs::remove_dir_all(directory).unwrap();
|
||
}
|
||
|
||
#[test]
|
||
fn risky_shell_commands_require_one_time_approval() {
|
||
let root = Path::new("/tmp/project");
|
||
assert!(risky_shell_reason("cargo test --all-features", root).is_none());
|
||
assert!(risky_shell_reason("git status --short", root).is_none());
|
||
for command in [
|
||
"rm -rf target",
|
||
"sudo make install",
|
||
"curl https://example.com",
|
||
"open report.html",
|
||
"touch ../outside",
|
||
"git push origin main",
|
||
] {
|
||
assert!(risky_shell_reason(command, root).is_some(), "{command}");
|
||
}
|
||
assert!(!SHELL_ENV_ALLOWLIST.contains(&"GITHUB_TOKEN"));
|
||
assert!(!SHELL_ENV_ALLOWLIST.contains(&"AWS_SECRET_ACCESS_KEY"));
|
||
assert!(!SHELL_ENV_ALLOWLIST.contains(&"SSH_AUTH_SOCK"));
|
||
}
|
||
|
||
#[test]
|
||
fn denial_and_stop_cancel_commands_awaiting_approval() {
|
||
let directory = std::env::temp_dir().join(format!(
|
||
"ds4-agent-approval-{}",
|
||
SystemTime::now()
|
||
.duration_since(UNIX_EPOCH)
|
||
.unwrap()
|
||
.as_nanos()
|
||
));
|
||
fs::create_dir_all(&directory).unwrap();
|
||
fs::write(directory.join("keep.txt"), "keep").unwrap();
|
||
|
||
let tools = Arc::new(Mutex::new(Tools::new(&directory, 4096).unwrap()));
|
||
let active = execute_async(
|
||
Arc::clone(&tools),
|
||
vec![call("bash", [("command", "rm keep.txt")])],
|
||
);
|
||
let decision = loop {
|
||
match active.events.recv_timeout(Duration::from_secs(2)).unwrap() {
|
||
ToolEvent::Approval {
|
||
prompt, decision, ..
|
||
} => {
|
||
assert!(prompt.detail.contains("rm keep.txt"));
|
||
assert_eq!(prompt.working_directory, directory.canonicalize().unwrap());
|
||
break decision;
|
||
}
|
||
ToolEvent::State { .. } => {}
|
||
}
|
||
};
|
||
decision.send(false).unwrap();
|
||
assert!(
|
||
active
|
||
.results
|
||
.recv_timeout(Duration::from_secs(2))
|
||
.unwrap()
|
||
.contains("user denied")
|
||
);
|
||
assert!(directory.join("keep.txt").exists());
|
||
|
||
let active = execute_async(tools, vec![call("google_search", [("query", "DS4")])]);
|
||
let _decision = loop {
|
||
match active.events.recv_timeout(Duration::from_secs(2)).unwrap() {
|
||
ToolEvent::Approval {
|
||
prompt, decision, ..
|
||
} => {
|
||
assert!(prompt.title.contains("browser"));
|
||
break decision;
|
||
}
|
||
ToolEvent::State { .. } => {}
|
||
}
|
||
};
|
||
active.cancel.store(true, Ordering::Relaxed);
|
||
assert!(
|
||
active
|
||
.results
|
||
.recv_timeout(Duration::from_secs(2))
|
||
.unwrap()
|
||
.contains("interrupted")
|
||
);
|
||
fs::remove_dir_all(directory).unwrap();
|
||
}
|
||
|
||
#[test]
|
||
fn oversized_shell_output_is_bounded_and_retained_until_cleanup() {
|
||
let directory = std::env::temp_dir().join(format!(
|
||
"ds4-agent-output-{}",
|
||
SystemTime::now()
|
||
.duration_since(UNIX_EPOCH)
|
||
.unwrap()
|
||
.as_nanos()
|
||
));
|
||
fs::create_dir_all(&directory).unwrap();
|
||
let mut tools = Tools::new(&directory, 4096).unwrap();
|
||
let result = tools.execute(
|
||
&call(
|
||
"bash",
|
||
[("command", "/usr/bin/yes x | /usr/bin/head -c 20000")],
|
||
),
|
||
&AtomicBool::new(false),
|
||
);
|
||
assert!(result.len() <= tools.result_limit());
|
||
assert!(result.contains("middle output omitted"));
|
||
let output = tool_output_path(&result).unwrap();
|
||
assert_eq!(fs::metadata(&output).unwrap().len(), 20_000);
|
||
fs::write(directory.join("large.txt"), "match xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\n".repeat(500)).unwrap();
|
||
let search = tools.execute(
|
||
&call("search", [("query", "match"), ("max_results", "500")]),
|
||
&AtomicBool::new(false),
|
||
);
|
||
assert!(search.contains("Use more to continue"));
|
||
assert!(
|
||
!tools
|
||
.execute(&call("more", []), &AtomicBool::new(false))
|
||
.is_empty()
|
||
);
|
||
drop(tools);
|
||
assert!(!output.exists());
|
||
fs::remove_dir_all(directory).unwrap();
|
||
}
|
||
|
||
#[test]
|
||
fn dropping_session_tools_stops_live_jobs_and_removes_output() {
|
||
let directory = std::env::temp_dir().join(format!(
|
||
"ds4-agent-switch-{}",
|
||
SystemTime::now()
|
||
.duration_since(UNIX_EPOCH)
|
||
.unwrap()
|
||
.as_nanos()
|
||
));
|
||
fs::create_dir_all(&directory).unwrap();
|
||
let mut tools = Tools::new(&directory, 4096).unwrap();
|
||
let result = tools.execute(
|
||
&call("bash", [("command", "sleep 30"), ("refresh_sec", "1")]),
|
||
&AtomicBool::new(false),
|
||
);
|
||
assert!(result.contains("status=running"));
|
||
let pid = result
|
||
.split_whitespace()
|
||
.find_map(|field| field.strip_prefix("pid="))
|
||
.unwrap()
|
||
.parse::<u32>()
|
||
.unwrap();
|
||
let output = tool_output_path(&result).unwrap();
|
||
drop(tools);
|
||
assert!(
|
||
!Command::new("/bin/kill")
|
||
.args(["-0", &pid.to_string()])
|
||
.status()
|
||
.unwrap()
|
||
.success()
|
||
);
|
||
assert!(!output.exists());
|
||
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 observation = tools.compaction_observation().unwrap();
|
||
assert!(observation.contains("bash job=2"));
|
||
assert!(observation.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(),
|
||
}
|
||
}
|
||
}
|