From c648b77ae4dccd8c81a94aced2dcc9a8dbd4b01e Mon Sep 17 00:00:00 2001 From: lixin Date: Thu, 4 Jun 2026 11:04:13 +0800 Subject: [PATCH 1/5] fix: render confirmation panel right border --- crates/aish-shell/src/app.rs | 116 ++++++++++++++++++++++------------- 1 file changed, 75 insertions(+), 41 deletions(-) diff --git a/crates/aish-shell/src/app.rs b/crates/aish-shell/src/app.rs index 108069e3..289935be 100644 --- a/crates/aish-shell/src/app.rs +++ b/crates/aish-shell/src/app.rs @@ -1088,31 +1088,38 @@ impl AishShell { .ok() .and_then(|s| s.parse::().ok()) .unwrap_or(80); - let border = "─".repeat(width.saturating_sub(4)); + let inner_width = width.saturating_sub(4).max(20); + let border = "─".repeat(inner_width); println!(); println!("\x1b[33m╭{}╮\x1b[0m", border); - println!( - "\x1b[33m│\x1b[1;33m ⚠ Security Confirmation Required\x1b[0m{}", - pad_to_width("", width.saturating_sub(38)) + print_panel_line( + "\x1b[1;33m ⚠ Security Confirmation Required\x1b[0m", + inner_width, ); - println!("\x1b[33m│\x1b[0m"); - println!( - "\x1b[33m│\x1b[0m \x1b[1;36m{}\x1b[0m {}", - t("shell.confirm_dialog_tool"), - ctx.tool_name + print_panel_line("", inner_width); + print_panel_line( + &format!( + " \x1b[1;36m{}\x1b[0m {}", + t("shell.confirm_dialog_tool"), + ctx.tool_name + ), + inner_width, ); let reason_lines = wrap_text(&ctx.message, width.saturating_sub(14)); - println!( - "\x1b[33m│\x1b[0m \x1b[1;36mReason:\x1b[0m {}", - reason_lines.lines().next().unwrap_or("") + print_panel_line( + &format!( + " \x1b[1;36mReason:\x1b[0m {}", + reason_lines.lines().next().unwrap_or("") + ), + inner_width, ); for line in reason_lines.lines().skip(1) { - println!("\x1b[33m│\x1b[0m {}", line); + print_panel_line(&format!(" {}", line), inner_width); } - println!("\x1b[33m│\x1b[0m"); - println!( - "\x1b[33m│\x1b[0m \x1b[36m{}\x1b[0m", - t("shell.confirm_dialog_question") + print_panel_line("", inner_width); + print_panel_line( + &format!(" \x1b[36m{}\x1b[0m", t("shell.confirm_dialog_question")), + inner_width, ); println!("\x1b[33m╰{}╯\x1b[0m", border); print!(" "); @@ -1130,28 +1137,36 @@ impl AishShell { .ok() .and_then(|s| s.parse::().ok()) .unwrap_or(80); - let border = "─".repeat(width.saturating_sub(4)); + let inner_width = width.saturating_sub(4).max(20); + let border = "─".repeat(inner_width); println!(); println!("\x1b[33m╭{}╮\x1b[0m", border); - println!( - "\x1b[33m│\x1b[1;33m {}\x1b[0m{}", - aish_i18n::t("shell.session.iteration_limit_title"), - pad_to_width("", width.saturating_sub(38)) + print_panel_line( + &format!( + "\x1b[1;33m {}\x1b[0m", + aish_i18n::t("shell.session.iteration_limit_title") + ), + inner_width, ); - println!("\x1b[33m│\x1b[0m"); - println!( - "\x1b[33m│\x1b[0m {} {}", - aish_i18n::t_with_args("shell.session.iteration_limit_reached", &{ - let mut m = std::collections::HashMap::new(); - m.insert("count".to_string(), iterations.to_string()); - m - },), - pad_to_width("", 0) + print_panel_line("", inner_width); + print_panel_line( + &format!( + " {}", + aish_i18n::t_with_args("shell.session.iteration_limit_reached", &{ + let mut m = std::collections::HashMap::new(); + m.insert("count".to_string(), iterations.to_string()); + m + },) + ), + inner_width, ); - println!("\x1b[33m│\x1b[0m"); - println!( - "\x1b[33m│\x1b[0m \x1b[36m{}\x1b[0m", - aish_i18n::t("shell.session.iteration_continue_prompt") + print_panel_line("", inner_width); + print_panel_line( + &format!( + " \x1b[36m{}\x1b[0m", + aish_i18n::t("shell.session.iteration_continue_prompt") + ), + inner_width, ); println!("\x1b[33m╰{}╯\x1b[0m", border); print!(" "); @@ -4976,13 +4991,32 @@ fn truncate_str(s: &str, max_len: usize) -> String { format!("{}...", truncated) } -/// Pad a string with trailing spaces to fill the given width (for box borders). -fn pad_to_width(s: &str, width: usize) -> String { - if s.len() >= width { - s.to_string() - } else { - format!("{}{}\x1b[33m│\x1b[0m", s, " ".repeat(width - s.len())) +fn print_panel_line(content: &str, inner_width: usize) { + let visible = ansi_display_width(content); + let padding = inner_width.saturating_sub(visible); + println!( + "\x1b[33m│\x1b[0m{}{}\x1b[33m│\x1b[0m", + content, + " ".repeat(padding) + ); +} + +fn ansi_display_width(s: &str) -> usize { + let mut width = 0usize; + let mut chars = s.chars().peekable(); + while let Some(ch) = chars.next() { + if ch == '\x1b' && chars.peek() == Some(&'[') { + let _ = chars.next(); + for code_ch in chars.by_ref() { + if code_ch.is_ascii_alphabetic() { + break; + } + } + continue; + } + width += unicode_width::UnicodeWidthChar::width(ch).unwrap_or(0); } + width } /// Wrap text to the given width, preserving word boundaries. From 2cf40c9d7292b05208f30595b1a1f47a2525dae9 Mon Sep 17 00:00:00 2001 From: lixin Date: Thu, 4 Jun 2026 20:16:40 +0800 Subject: [PATCH 2/5] feat: add web_fetch tool --- Cargo.lock | 3 + README.md | 2 +- README_CN.md | 2 +- crates/aish-i18n/locales/en-US.yaml | 11 + crates/aish-i18n/locales/zh-CN.yaml | 11 + crates/aish-shell/src/app.rs | 7 + crates/aish-tools/Cargo.toml | 3 + crates/aish-tools/src/lib.rs | 2 + crates/aish-tools/src/web_fetch.rs | 935 ++++++++++++++++++++++++++++ 9 files changed, 974 insertions(+), 2 deletions(-) create mode 100644 crates/aish-tools/src/web_fetch.rs diff --git a/Cargo.lock b/Cargo.lock index 1005a896..e67ec046 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -292,13 +292,16 @@ dependencies = [ "aish-pty", "aish-security", "aish-ui", + "futures", "glob", "inquire", "regex", + "reqwest", "serde", "serde_json", "shellexpand", "tempfile", + "tokio", "tracing", "uuid", ] diff --git a/README.md b/README.md index 0925cd04..ff001fd2 100644 --- a/README.md +++ b/README.md @@ -368,7 +368,7 @@ cargo clippy --all-targets -- -D warnings | `aish-security` | Security policy engine with glob-to-regex pattern matching | | `aish-skills` | Skill plugin discovery with hot-reload via notify | | `aish-memory` | Markdown-based long-term memory with relevance scoring | -| `aish-tools` | Built-in tools: bash, fs (read/write/edit), ask_user, memory, skill | +| `aish-tools` | Built-in tools: bash, WebFetch, fs (read/write/edit), ask_user, memory, skill | | `aish-scripts` | .aish script system with frontmatter, ai "prompt" syntax, hooks | | `aish-shell` | Main shell: REPL loop, AI handler, built-in commands, animation | | `aish-cli` | CLI entry point with clap derive macros | diff --git a/README_CN.md b/README_CN.md index e88b5798..2c26d048 100644 --- a/README_CN.md +++ b/README_CN.md @@ -370,7 +370,7 @@ cargo clippy | `aish-security` | 安全策略引擎,glob-to-regex 模式匹配 | | `aish-skills` | 技能插件发现,notify 热加载 | | `aish-memory` | Markdown 长期记忆,相关性评分 | -| `aish-tools` | 内置工具:bash、文件读写编辑、ask_user、memory、skill | +| `aish-tools` | 内置工具:bash、WebFetch、文件读写编辑、ask_user、memory、skill | | `aish-scripts` | .aish 脚本系统,frontmatter、ai "prompt" 语法、hooks | | `aish-shell` | 主 Shell:REPL 循环、AI 处理器、内置命令、动画 | | `aish-cli` | CLI 入口,clap derive 宏 | diff --git a/crates/aish-i18n/locales/en-US.yaml b/crates/aish-i18n/locales/en-US.yaml index 4741bd0c..e3ea1645 100644 --- a/crates/aish-i18n/locales/en-US.yaml +++ b/crates/aish-i18n/locales/en-US.yaml @@ -806,6 +806,17 @@ tools: missing_pattern: "Missing 'pattern' parameter" invalid_glob: "Error: invalid glob pattern: {error}" + web_fetch: + description: "Fetch content from a specified URL, convert readable HTML to text, and answer a prompt about the page using a secondary model. IMPORTANT: WebFetch will fail for authenticated or private URLs. For GitHub URLs, prefer gh via bash when available." + missing_url: "Missing 'url' parameter" + missing_prompt: "Missing 'prompt' parameter" + invalid_url: "Error: invalid URL. Provide a fully-qualified http or https URL without credentials." + blocked_private_host: "Blocked WebFetch request to private or local host: {host}" + confirm_fetch: "Allow WebFetch to fetch content from {host}?" + content_too_large: "WebFetch response exceeded the 10MB limit" + binary_unsupported: "WebFetch does not support binary content in this version" + secondary_model_failed: "WebFetch fetched the page but failed to process it with the secondary model: {error}" + help: labels: usage: "Usage" diff --git a/crates/aish-i18n/locales/zh-CN.yaml b/crates/aish-i18n/locales/zh-CN.yaml index e5431c6e..c79e25a8 100644 --- a/crates/aish-i18n/locales/zh-CN.yaml +++ b/crates/aish-i18n/locales/zh-CN.yaml @@ -805,6 +805,17 @@ tools: missing_pattern: "缺少 'pattern' 参数" invalid_glob: "错误:无效的 glob 模式: {error}" + web_fetch: + description: "从指定 URL 抓取内容,将可读 HTML 转为文本,并使用二级模型按 prompt 分析页面。重要:WebFetch 无法访问需要认证或私有的 URL。对于 GitHub URL,如可用请优先通过 bash 使用 gh。" + missing_url: "缺少 'url' 参数" + missing_prompt: "缺少 'prompt' 参数" + invalid_url: "错误:无效 URL。请提供完整的 http 或 https URL,且不要包含凭据。" + blocked_private_host: "已阻止访问私有或本地主机的 WebFetch 请求: {host}" + confirm_fetch: "允许 WebFetch 从 {host} 抓取内容吗?" + content_too_large: "WebFetch 响应超过 10MB 限制" + binary_unsupported: "当前版本的 WebFetch 不支持二进制内容" + secondary_model_failed: "WebFetch 已抓取页面,但使用二级模型处理失败: {error}" + help: labels: usage: "用法" diff --git a/crates/aish-shell/src/app.rs b/crates/aish-shell/src/app.rs index 289935be..9d532045 100644 --- a/crates/aish-shell/src/app.rs +++ b/crates/aish-shell/src/app.rs @@ -445,6 +445,13 @@ impl AishShell { tool_registry.register(Box::new(aish_tools::PythonTool::new())); tool_registry.register(Box::new(aish_tools::GlobTool::new())); tool_registry.register(Box::new(aish_tools::GrepTool::new())); + tool_registry.register(Box::new(aish_tools::WebFetchTool::new( + &config.api_base, + &config.api_key, + &config.model, + Some(config.temperature), + config.max_tokens, + ))); tool_registry.register(Box::new(aish_tools::EnterPlanModeTool::new())); tool_registry.register(Box::new(aish_tools::ExitPlanModeTool::new())); diff --git a/crates/aish-tools/Cargo.toml b/crates/aish-tools/Cargo.toml index 1bfc78a7..aea0a148 100644 --- a/crates/aish-tools/Cargo.toml +++ b/crates/aish-tools/Cargo.toml @@ -15,6 +15,9 @@ serde_json.workspace = true tracing.workspace = true regex.workspace = true uuid.workspace = true +reqwest.workspace = true +tokio.workspace = true +futures.workspace = true glob = "0.3" shellexpand = "2.1" inquire = "0.9.4" diff --git a/crates/aish-tools/src/lib.rs b/crates/aish-tools/src/lib.rs index f6ce4547..4eacf716 100644 --- a/crates/aish-tools/src/lib.rs +++ b/crates/aish-tools/src/lib.rs @@ -29,6 +29,7 @@ pub mod registry; pub mod secure_bash; pub mod skill_tool; pub mod system_diagnose; +pub mod web_fetch; pub use ask_user::AskUserTool; pub use channel_ask_user::ChannelAskUserTool; @@ -46,3 +47,4 @@ pub use secure_bash::SecureBashTool; pub use skill_tool::{SkillInfo, SkillTool}; pub use system_diagnose::SharedEventCallback; pub use system_diagnose::SystemDiagnoseTool; +pub use web_fetch::WebFetchTool; diff --git a/crates/aish-tools/src/web_fetch.rs b/crates/aish-tools/src/web_fetch.rs new file mode 100644 index 00000000..18436b6d --- /dev/null +++ b/crates/aish-tools/src/web_fetch.rs @@ -0,0 +1,935 @@ +use std::collections::HashMap; +use std::future::Future; +use std::net::IpAddr; +use std::pin::Pin; +use std::sync::{Mutex, OnceLock}; +use std::time::{Duration, Instant}; + +use aish_llm::{ + ChatMessage, LlmResponse, LlmSession, PreflightResult, PreflightSecurityContext, + SecurityPanelMode, StreamParser, Tool, ToolResult, +}; +use futures::StreamExt; +use regex::Regex; +use reqwest::header::{ACCEPT, CONTENT_TYPE, USER_AGENT}; +use reqwest::{redirect, Client, StatusCode, Url}; + +const TOOL_NAME: &str = "WebFetch"; +const MAX_URL_LENGTH: usize = 2000; +const MAX_HTTP_CONTENT_LENGTH: usize = 10 * 1024 * 1024; +const FETCH_TIMEOUT_SECS: u64 = 60; +const MAX_REDIRECTS: usize = 10; +const MAX_MARKDOWN_LENGTH: usize = 100_000; +const CACHE_TTL: Duration = Duration::from_secs(15 * 60); +const CACHE_MAX_ENTRIES: usize = 64; +const USER_AGENT_VALUE: &str = concat!("aish/", env!("CARGO_PKG_VERSION"), " WebFetch"); + +#[derive(Clone)] +struct CacheEntry { + fetched_at: Instant, + url: String, + code: u16, + code_text: String, + bytes: usize, + content_type: String, + content: String, +} + +static URL_CACHE: OnceLock>> = OnceLock::new(); +static DESCRIPTION: OnceLock = OnceLock::new(); + +fn cache() -> &'static Mutex> { + URL_CACHE.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn get_description() -> &'static str { + DESCRIPTION.get_or_init(|| aish_i18n::t("tools.web_fetch.description")) +} + +/// Fetch a URL, extract readable text, and answer a focused prompt about it. +pub struct WebFetchTool { + api_base: String, + api_key: String, + model: String, + temperature: Option, + max_tokens: Option, +} + +impl WebFetchTool { + pub fn new( + api_base: &str, + api_key: &str, + model: &str, + temperature: Option, + max_tokens: Option, + ) -> Self { + Self { + api_base: api_base.to_string(), + api_key: api_key.to_string(), + model: model.to_string(), + temperature, + max_tokens, + } + } + + fn build_client() -> Result { + Client::builder() + .redirect(redirect::Policy::none()) + .timeout(Duration::from_secs(FETCH_TIMEOUT_SECS)) + .build() + .map_err(|error| error.to_string()) + } + + async fn fetch_url_content(&self, raw_url: &str) -> Result { + let normalized = validate_and_normalize_url(raw_url).map_err(FetchFailure::Blocked)?; + ensure_public_host(&normalized) + .await + .map_err(FetchFailure::Blocked)?; + + if let Some(entry) = get_cached(raw_url) { + return Ok(FetchedContent { + url: entry.url, + code: entry.code, + code_text: entry.code_text, + bytes: entry.bytes, + content_type: entry.content_type, + content: entry.content, + from_cache: true, + }); + } + + let client = Self::build_client().map_err(FetchFailure::Request)?; + let response = get_with_permitted_redirects(&client, normalized.clone(), 0).await?; + let status = response.status(); + let content_type = response + .headers() + .get(CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or("") + .to_string(); + + if is_binary_content_type(&content_type) { + return Err(FetchFailure::Request(aish_i18n::t( + "tools.web_fetch.binary_unsupported", + ))); + } + + let raw = read_limited_body(response) + .await + .map_err(FetchFailure::Request)?; + let bytes = raw.len(); + let text = String::from_utf8_lossy(&raw).to_string(); + let content = if content_type.to_ascii_lowercase().contains("text/html") { + html_to_readable_text(&text) + } else { + text + }; + + let entry = CacheEntry { + fetched_at: Instant::now(), + url: normalized.to_string(), + code: status.as_u16(), + code_text: status_text(status).to_string(), + bytes, + content_type: content_type.clone(), + content: content.clone(), + }; + set_cached(raw_url.to_string(), entry); + + Ok(FetchedContent { + url: normalized.to_string(), + code: status.as_u16(), + code_text: status_text(status).to_string(), + bytes, + content_type, + content, + from_cache: false, + }) + } + + async fn apply_prompt_to_content( + &self, + prompt: &str, + content: &str, + is_preapproved_domain: bool, + ) -> Result { + let model_prompt = make_secondary_model_prompt(content, prompt, is_preapproved_domain); + let session = LlmSession::new( + &self.api_base, + &self.api_key, + &self.model, + self.temperature.or(Some(0.1)), + self.max_tokens.or(Some(2048)), + ); + let messages = vec![ChatMessage::system(""), ChatMessage::user(model_prompt)]; + match session + .chat_completion_raw(&messages, None, false, Some(0.1), Some(2048)) + .await + .map_err(|error| error.to_string())? + { + LlmResponse::Json(json) => { + let (content, _reasoning, _tool_calls, _usage) = + StreamParser::parse_response(&json); + Ok(content.unwrap_or_else(|| "No response from model".to_string())) + } + LlmResponse::Stream(_) => Err("secondary model unexpectedly returned a stream".into()), + } + } +} + +impl Tool for WebFetchTool { + fn name(&self) -> &str { + TOOL_NAME + } + + fn description(&self) -> &str { + get_description() + } + + fn parameters(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "url": { + "type": "string", + "description": "The fully-qualified URL to fetch content from" + }, + "prompt": { + "type": "string", + "description": "The prompt describing what information to extract from the fetched page" + } + }, + "required": ["url", "prompt"], + "additionalProperties": false + }) + } + + fn preflight(&self, args: &serde_json::Value) -> PreflightResult { + let url = match args.get("url").and_then(|value| value.as_str()) { + Some(value) if !value.trim().is_empty() => value, + _ => { + return PreflightResult::Block { + message: aish_i18n::t("tools.web_fetch.missing_url"), + security: Some(PreflightSecurityContext::fallback( + TOOL_NAME, + None, + aish_i18n::t("tools.web_fetch.missing_url"), + SecurityPanelMode::Blocked, + )), + } + } + }; + + let normalized = match validate_and_normalize_url(url) { + Ok(parsed) => parsed, + Err(message) => { + return PreflightResult::Block { + message: message.clone(), + security: Some(PreflightSecurityContext::fallback( + TOOL_NAME, + Some(url.to_string()), + message, + SecurityPanelMode::Blocked, + )), + } + } + }; + + let hostname = normalized.host_str().unwrap_or("").to_string(); + if is_preapproved_host(&hostname, normalized.path()) { + return PreflightResult::Allow; + } + + let message = aish_i18n::t_with_args( + "tools.web_fetch.confirm_fetch", + &HashMap::from([("host".to_string(), hostname.clone())]), + ); + PreflightResult::Confirm { + message: message.clone(), + security: Some(PreflightSecurityContext::fallback( + TOOL_NAME, + Some(hostname), + message, + SecurityPanelMode::Confirm, + )), + } + } + + fn execute(&self, _args: serde_json::Value) -> ToolResult { + ToolResult::error("WebFetch requires async execution; use execute_async") + } + + fn execute_async<'a>( + &'a self, + args: serde_json::Value, + ) -> Pin + Send + 'a>> { + Box::pin(async move { + let url = match args.get("url").and_then(|value| value.as_str()) { + Some(value) if !value.trim().is_empty() => value.trim(), + _ => return ToolResult::error(aish_i18n::t("tools.web_fetch.missing_url")), + }; + let prompt = match args.get("prompt").and_then(|value| value.as_str()) { + Some(value) if !value.trim().is_empty() => value.trim(), + _ => return ToolResult::error(aish_i18n::t("tools.web_fetch.missing_prompt")), + }; + + let start = Instant::now(); + let fetched = match self.fetch_url_content(url).await { + Ok(content) => content, + Err(FetchFailure::Redirect(info)) => { + let message = format_redirect_message(&info, prompt); + return ToolResult { + ok: true, + output: message.clone(), + meta: Some(serde_json::json!({ + "url": url, + "redirect_url": info.redirect_url, + "code": info.status_code, + "result": message, + "durationMs": start.elapsed().as_millis() as u64, + })), + }; + } + Err(FetchFailure::Blocked(message)) | Err(FetchFailure::Request(message)) => { + return ToolResult::error(message) + } + }; + + let truncated_content = truncate_for_model(&fetched.content); + let parsed_url = Url::parse(&fetched.url).ok(); + let is_preapproved_domain = parsed_url + .as_ref() + .and_then(|parsed| parsed.host_str().map(|host| (host, parsed.path()))) + .is_some_and(|(host, path)| is_preapproved_host(host, path)); + + let result = match self + .apply_prompt_to_content(prompt, &truncated_content, is_preapproved_domain) + .await + { + Ok(result) => result, + Err(error) => { + let mut args_map = HashMap::new(); + args_map.insert("error".to_string(), error); + return ToolResult::error(aish_i18n::t_with_args( + "tools.web_fetch.secondary_model_failed", + &args_map, + )); + } + }; + + let duration_ms = start.elapsed().as_millis() as u64; + let output = format!( + "Fetched: {}\nStatus: {} {}\nBytes: {}\nDuration: {}ms\nCached: {}\n\n{}", + fetched.url, + fetched.code, + fetched.code_text, + fetched.bytes, + duration_ms, + fetched.from_cache, + result + ); + + ToolResult { + ok: true, + output, + meta: Some(serde_json::json!({ + "url": fetched.url, + "code": fetched.code, + "codeText": fetched.code_text, + "bytes": fetched.bytes, + "contentType": fetched.content_type, + "durationMs": duration_ms, + "fromCache": fetched.from_cache, + "result": result, + })), + } + }) + } +} + +#[derive(Debug)] +enum FetchFailure { + Blocked(String), + Request(String), + Redirect(RedirectInfo), +} + +#[derive(Debug)] +struct RedirectInfo { + original_url: String, + redirect_url: String, + status_code: u16, +} + +struct FetchedContent { + url: String, + code: u16, + code_text: String, + bytes: usize, + content_type: String, + content: String, + from_cache: bool, +} + +async fn get_with_permitted_redirects( + client: &Client, + url: Url, + depth: usize, +) -> Result { + if depth > MAX_REDIRECTS { + return Err(FetchFailure::Request(format!( + "Too many redirects (exceeded {})", + MAX_REDIRECTS + ))); + } + + ensure_public_host(&url) + .await + .map_err(FetchFailure::Blocked)?; + let response = client + .get(url.clone()) + .header(ACCEPT, "text/markdown, text/html, text/plain, */*") + .header(USER_AGENT, USER_AGENT_VALUE) + .send() + .await + .map_err(|error| FetchFailure::Request(error.to_string()))?; + + if is_redirect_status(response.status()) { + let location = response + .headers() + .get(reqwest::header::LOCATION) + .and_then(|value| value.to_str().ok()) + .ok_or_else(|| FetchFailure::Request("Redirect missing Location header".into()))?; + let redirect_url = url + .join(location) + .map_err(|error| FetchFailure::Request(error.to_string()))?; + validate_url_basics(&redirect_url).map_err(FetchFailure::Blocked)?; + ensure_public_host(&redirect_url) + .await + .map_err(FetchFailure::Blocked)?; + + if is_permitted_redirect(&url, &redirect_url) { + return Box::pin(get_with_permitted_redirects( + client, + redirect_url, + depth + 1, + )) + .await; + } + + return Err(FetchFailure::Redirect(RedirectInfo { + original_url: url.to_string(), + redirect_url: redirect_url.to_string(), + status_code: response.status().as_u16(), + })); + } + + Ok(response) +} + +fn validate_and_normalize_url(raw_url: &str) -> Result { + if raw_url.len() > MAX_URL_LENGTH { + return Err(aish_i18n::t("tools.web_fetch.invalid_url")); + } + let mut parsed = + Url::parse(raw_url).map_err(|_| aish_i18n::t("tools.web_fetch.invalid_url"))?; + if parsed.scheme() == "http" { + parsed + .set_scheme("https") + .map_err(|_| aish_i18n::t("tools.web_fetch.invalid_url"))?; + } + validate_url_basics(&parsed)?; + Ok(parsed) +} + +fn validate_url_basics(parsed: &Url) -> Result<(), String> { + if parsed.scheme() != "https" && parsed.scheme() != "http" { + return Err(aish_i18n::t("tools.web_fetch.invalid_url")); + } + if !parsed.username().is_empty() || parsed.password().is_some() { + return Err(aish_i18n::t("tools.web_fetch.invalid_url")); + } + let host = parsed + .host_str() + .ok_or_else(|| aish_i18n::t("tools.web_fetch.invalid_url"))?; + if host.split('.').count() < 2 && host.parse::().is_err() { + return Err(aish_i18n::t("tools.web_fetch.invalid_url")); + } + if is_blocked_hostname(host) { + let mut args = HashMap::new(); + args.insert("host".to_string(), host.to_string()); + return Err(aish_i18n::t_with_args( + "tools.web_fetch.blocked_private_host", + &args, + )); + } + Ok(()) +} + +async fn ensure_public_host(url: &Url) -> Result<(), String> { + let host = url + .host_str() + .ok_or_else(|| aish_i18n::t("tools.web_fetch.invalid_url"))?; + if is_blocked_hostname(host) { + let mut args = HashMap::new(); + args.insert("host".to_string(), host.to_string()); + return Err(aish_i18n::t_with_args( + "tools.web_fetch.blocked_private_host", + &args, + )); + } + if host.parse::().is_ok() { + return Ok(()); + } + + let port = url.port_or_known_default().unwrap_or(443); + let addrs = tokio::net::lookup_host((host, port)) + .await + .map_err(|error| format!("DNS lookup failed for {}: {}", host, error))?; + let mut found = false; + for addr in addrs { + found = true; + if is_private_ip(&addr.ip()) { + let mut args = HashMap::new(); + args.insert("host".to_string(), host.to_string()); + return Err(aish_i18n::t_with_args( + "tools.web_fetch.blocked_private_host", + &args, + )); + } + } + if !found { + return Err(format!("DNS lookup returned no addresses for {}", host)); + } + Ok(()) +} + +fn is_blocked_hostname(host: &str) -> bool { + let normalized = host.trim_end_matches('.').to_ascii_lowercase(); + if matches!( + normalized.as_str(), + "localhost" | "metadata.google.internal" + ) { + return true; + } + if normalized.ends_with(".localhost") || normalized.ends_with(".local") { + return true; + } + match normalized.parse::() { + Ok(ip) => is_private_ip(&ip), + Err(_) => false, + } +} + +fn is_private_ip(ip: &IpAddr) -> bool { + match ip { + IpAddr::V4(addr) => { + addr.is_private() + || addr.is_loopback() + || addr.is_link_local() + || addr.is_broadcast() + || addr.is_documentation() + || addr.is_unspecified() + || addr.octets() == [169, 254, 169, 254] + || addr.octets()[0] == 0 + } + IpAddr::V6(addr) => { + addr.is_loopback() + || addr.is_unspecified() + || addr.is_unique_local() + || addr.is_unicast_link_local() + } + } +} + +fn is_redirect_status(status: StatusCode) -> bool { + matches!( + status, + StatusCode::MOVED_PERMANENTLY + | StatusCode::FOUND + | StatusCode::TEMPORARY_REDIRECT + | StatusCode::PERMANENT_REDIRECT + ) +} + +fn is_permitted_redirect(original: &Url, redirect_url: &Url) -> bool { + if original.scheme() != redirect_url.scheme() || original.port() != redirect_url.port() { + return false; + } + if !redirect_url.username().is_empty() || redirect_url.password().is_some() { + return false; + } + let Some(original_host) = original.host_str() else { + return false; + }; + let Some(redirect_host) = redirect_url.host_str() else { + return false; + }; + strip_www(original_host) == strip_www(redirect_host) +} + +fn strip_www(host: &str) -> &str { + host.strip_prefix("www.").unwrap_or(host) +} + +async fn read_limited_body(response: reqwest::Response) -> Result, String> { + if response + .content_length() + .is_some_and(|length| length > MAX_HTTP_CONTENT_LENGTH as u64) + { + return Err(aish_i18n::t("tools.web_fetch.content_too_large")); + } + + let mut stream = response.bytes_stream(); + let mut buffer = Vec::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|error| error.to_string())?; + if buffer.len() + chunk.len() > MAX_HTTP_CONTENT_LENGTH { + return Err(aish_i18n::t("tools.web_fetch.content_too_large")); + } + buffer.extend_from_slice(&chunk); + } + Ok(buffer) +} + +fn status_text(status: StatusCode) -> &'static str { + status.canonical_reason().unwrap_or("") +} + +fn is_binary_content_type(content_type: &str) -> bool { + let lower = content_type.to_ascii_lowercase(); + if lower.starts_with("text/") { + return false; + } + if lower.contains("json") || lower.contains("xml") || lower.contains("javascript") { + return false; + } + lower.contains("application/pdf") + || lower.starts_with("image/") + || lower.starts_with("audio/") + || lower.starts_with("video/") + || lower.contains("application/octet-stream") +} + +fn html_to_readable_text(html: &str) -> String { + let without_scripts = regex_replace_all( + html, + r"(?is)]*>.*?|]*>.*?|]*>.*?|]*>.*?|]*>.*?", + "\n", + ); + let with_breaks = regex_replace_all( + &without_scripts, + r"(?i)]*>", + "\n", + ); + let without_tags = regex_replace_all(&with_breaks, r"(?is)<[^>]+>", " "); + normalize_text_whitespace(&decode_html_entities(&without_tags)) +} + +fn regex_replace_all(input: &str, pattern: &str, replacement: &str) -> String { + match Regex::new(pattern) { + Ok(regex) => regex.replace_all(input, replacement).to_string(), + Err(_) => input.to_string(), + } +} + +fn decode_html_entities(input: &str) -> String { + input + .replace(" ", " ") + .replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace(""", "\"") + .replace("'", "'") + .replace("'", "'") +} + +fn normalize_text_whitespace(input: &str) -> String { + let mut lines = Vec::new(); + for line in input.lines() { + let collapsed = line.split_whitespace().collect::>().join(" "); + if !collapsed.is_empty() { + lines.push(collapsed); + } + } + lines.join("\n") +} + +fn truncate_for_model(content: &str) -> String { + if content.chars().count() <= MAX_MARKDOWN_LENGTH { + return content.to_string(); + } + let mut truncated = content + .chars() + .take(MAX_MARKDOWN_LENGTH) + .collect::(); + truncated.push_str("\n\n[Content truncated due to length...]"); + truncated +} + +fn make_secondary_model_prompt( + markdown_content: &str, + prompt: &str, + is_preapproved_domain: bool, +) -> String { + let guidelines = if is_preapproved_domain { + "Provide a concise response based on the content above. Include relevant details, code examples, and documentation excerpts as needed." + } else { + "Provide a concise response based only on the content above. In your response:\n - Enforce a strict 125-character maximum for quotes from any source document. Open Source Software is ok as long as we respect the license.\n - Use quotation marks for exact language from articles; any language outside of the quotation should never be word-for-word the same.\n - You are not a lawyer and never comment on the legality of your own prompts and responses.\n - Never produce or reproduce exact song lyrics." + }; + format!( + "Web page content:\n---\n{}\n---\n\n{}\n\n{}\n", + markdown_content, prompt, guidelines + ) +} + +fn format_redirect_message(info: &RedirectInfo, prompt: &str) -> String { + let status_text = match info.status_code { + 301 => "Moved Permanently", + 307 => "Temporary Redirect", + 308 => "Permanent Redirect", + _ => "Found", + }; + format!( + "REDIRECT DETECTED: The URL redirects to a different host.\n\nOriginal URL: {}\nRedirect URL: {}\nStatus: {} {}\n\nTo complete your request, fetch the redirected URL with these parameters:\n- url: \"{}\"\n- prompt: \"{}\"", + info.original_url, + info.redirect_url, + info.status_code, + status_text, + info.redirect_url, + prompt + ) +} + +fn get_cached(key: &str) -> Option { + let mut guard = cache().lock().ok()?; + let now = Instant::now(); + guard.retain(|_, entry| now.duration_since(entry.fetched_at) < CACHE_TTL); + guard.get(key).cloned() +} + +fn set_cached(key: String, entry: CacheEntry) { + if let Ok(mut guard) = cache().lock() { + if guard.len() >= CACHE_MAX_ENTRIES { + if let Some(oldest_key) = guard + .iter() + .min_by_key(|(_, value)| value.fetched_at) + .map(|(cache_key, _)| cache_key.clone()) + { + guard.remove(&oldest_key); + } + } + guard.insert(key, entry); + } +} + +fn is_preapproved_host(hostname: &str, pathname: &str) -> bool { + for entry in PREAPPROVED_HOSTS { + if let Some((host, prefix)) = entry.split_once('/') { + if hostname == host { + let prefix = format!("/{}", prefix); + if pathname == prefix || pathname.starts_with(&(prefix + "/")) { + return true; + } + } + continue; + } + if hostname == *entry { + return true; + } + } + false +} + +const PREAPPROVED_HOSTS: &[&str] = &[ + "platform.claude.com", + "code.claude.com", + "modelcontextprotocol.io", + "github.com/anthropics", + "agentskills.io", + "docs.python.org", + "en.cppreference.com", + "docs.oracle.com", + "learn.microsoft.com", + "developer.mozilla.org", + "go.dev", + "pkg.go.dev", + "www.php.net", + "docs.swift.org", + "kotlinlang.org", + "ruby-doc.org", + "doc.rust-lang.org", + "www.typescriptlang.org", + "react.dev", + "angular.io", + "vuejs.org", + "nextjs.org", + "expressjs.com", + "nodejs.org", + "bun.sh", + "jquery.com", + "getbootstrap.com", + "tailwindcss.com", + "d3js.org", + "threejs.org", + "redux.js.org", + "webpack.js.org", + "jestjs.io", + "reactrouter.com", + "docs.djangoproject.com", + "flask.palletsprojects.com", + "fastapi.tiangolo.com", + "pandas.pydata.org", + "numpy.org", + "www.tensorflow.org", + "pytorch.org", + "scikit-learn.org", + "matplotlib.org", + "requests.readthedocs.io", + "jupyter.org", + "laravel.com", + "symfony.com", + "wordpress.org", + "docs.spring.io", + "hibernate.org", + "tomcat.apache.org", + "gradle.org", + "maven.apache.org", + "asp.net", + "dotnet.microsoft.com", + "nuget.org", + "blazor.net", + "reactnative.dev", + "docs.flutter.dev", + "developer.apple.com", + "developer.android.com", + "keras.io", + "spark.apache.org", + "huggingface.co", + "www.kaggle.com", + "www.mongodb.com", + "redis.io", + "www.postgresql.org", + "dev.mysql.com", + "www.sqlite.org", + "graphql.org", + "prisma.io", + "docs.aws.amazon.com", + "cloud.google.com", + "kubernetes.io", + "www.docker.com", + "www.terraform.io", + "www.ansible.com", + "vercel.com/docs", + "docs.netlify.com", + "devcenter.heroku.com", + "cypress.io", + "selenium.dev", + "docs.unity.com", + "docs.unrealengine.com", + "git-scm.com", + "nginx.org", + "httpd.apache.org", +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn normalizes_http_to_https() { + let url = validate_and_normalize_url("http://example.com/path").unwrap(); + assert_eq!(url.as_str(), "https://example.com/path"); + } + + #[test] + fn rejects_private_hosts() { + assert!(validate_and_normalize_url("https://localhost/").is_err()); + assert!(validate_and_normalize_url("https://127.0.0.1/").is_err()); + assert!(validate_and_normalize_url("https://169.254.169.254/").is_err()); + assert!(validate_and_normalize_url("https://10.0.0.5/").is_err()); + } + + #[test] + fn preapproved_host_supports_path_prefix_boundary() { + assert!(is_preapproved_host("github.com", "/anthropics/claude-code")); + assert!(!is_preapproved_host( + "github.com", + "/anthropics-evil/project" + )); + assert!(is_preapproved_host("doc.rust-lang.org", "/book/")); + } + + #[test] + fn redirect_only_allows_same_origin_or_www_equivalent() { + let original = Url::parse("https://example.com/docs").unwrap(); + let same = Url::parse("https://www.example.com/docs").unwrap(); + let other = Url::parse("https://evil.example.net/docs").unwrap(); + let http = Url::parse("http://example.com/docs").unwrap(); + assert!(is_permitted_redirect(&original, &same)); + assert!(!is_permitted_redirect(&original, &other)); + assert!(!is_permitted_redirect(&original, &http)); + } + + #[test] + fn html_to_readable_text_removes_scripts_and_tags() { + let html = "

Hello & hi

World

"; + let text = html_to_readable_text(html); + assert!(text.contains("Hello & hi")); + assert!(text.contains("World")); + assert!(!text.contains("bad()")); + assert!(!text.contains("

")); + } + + #[test] + fn secondary_prompt_includes_quote_restriction_for_unapproved_domains() { + let prompt = make_secondary_model_prompt("content", "summarize", false); + assert!(prompt.contains("125-character maximum")); + assert!(prompt.contains("summarize")); + } + + #[tokio::test] + #[ignore] + async fn live_fetch_url() { + if std::env::var("AISH_LIVE_WEBFETCH").ok().as_deref() != Some("1") { + eprintln!("set AISH_LIVE_WEBFETCH=1 to run this live network smoke test"); + return; + } + + let url = std::env::var("AISH_LIVE_WEBFETCH_URL") + .unwrap_or_else(|_| "https://github.com/mattpocock/skills".to_string()); + let expected = std::env::var("AISH_LIVE_WEBFETCH_EXPECT").ok(); + let tool = WebFetchTool::new("", "", "", Some(0.1), Some(256)); + let fetched = tool + .fetch_url_content(&url) + .await + .expect("expected live page fetch to succeed"); + + println!( + "fetched {} status={} bytes={} content_type={} chars={}", + fetched.url, + fetched.code, + fetched.bytes, + fetched.content_type, + fetched.content.len() + ); + println!( + "preview:\n{}", + truncate_for_model(&fetched.content) + .chars() + .take(800) + .collect::() + ); + + assert_eq!(fetched.code, 200); + if let Some(expected) = expected { + assert!( + fetched + .content + .to_ascii_lowercase() + .contains(&expected.to_ascii_lowercase()), + "expected fetched content to contain {expected:?}" + ); + } + } +} From beadbcd03c36c0e4e9d70ffb521a40a288a77cc7 Mon Sep 17 00:00:00 2001 From: lixin Date: Fri, 5 Jun 2026 14:28:11 +0800 Subject: [PATCH 3/5] feat(tool): refactor aish-tool crate - remove tool description from i18n crate - add webfetch tool --- crates/aish-i18n/locales/de-DE.yaml | 20 - crates/aish-i18n/locales/en-US.yaml | 68 -- crates/aish-i18n/locales/es-ES.yaml | 20 - crates/aish-i18n/locales/fr-FR.yaml | 20 - crates/aish-i18n/locales/ja-JP.yaml | 20 - crates/aish-i18n/locales/zh-CN.yaml | 68 -- crates/aish-llm/src/agent.rs | 3 +- crates/aish-llm/src/session.rs | 119 ++- crates/aish-llm/src/types.rs | 4 + .../aish-tools/src/{ => ask_user}/ask_user.rs | 82 +- crates/aish-tools/src/ask_user/prompt.rs | 81 ++ crates/aish-tools/src/{ => bash}/bash.rs | 31 +- crates/aish-tools/src/bash/prompt.rs | 27 + .../channel_ask_user.rs | 60 +- .../aish-tools/src/channel_ask_user/prompt.rs | 57 ++ .../src/{ => channel_bash}/channel_bash.rs | 28 +- crates/aish-tools/src/channel_bash/prompt.rs | 26 + crates/aish-tools/src/edit_file/edit_file.rs | 173 ++++ crates/aish-tools/src/edit_file/prompt.rs | 33 + .../src/{ => final_answer}/final_answer.rs | 19 +- crates/aish-tools/src/final_answer/prompt.rs | 20 + crates/aish-tools/src/fs.rs | 551 ----------- .../src/{ => glob_tool}/glob_tool.rs | 28 +- crates/aish-tools/src/glob_tool/prompt.rs | 25 + .../src/{ => grep_tool}/grep_tool.rs | 32 +- crates/aish-tools/src/grep_tool/prompt.rs | 30 + .../src/{ => host_note}/host_note.rs | 32 +- crates/aish-tools/src/host_note/prompt.rs | 30 + crates/aish-tools/src/lib.rs | 153 ++- .../src/{ => memory_tool}/memory_tool.rs | 42 +- crates/aish-tools/src/memory_tool/prompt.rs | 40 + crates/aish-tools/src/plan_tool.rs | 488 --------- .../src/plan_tool/enter_plan_mode.rs | 191 ++++ .../src/plan_tool/exit_plan_mode.rs | 155 +++ .../src/plan_tool/list_plan_templates.rs | 100 ++ crates/aish-tools/src/plan_tool/prompt.rs | 63 ++ crates/aish-tools/src/python/prompt.rs | 22 + crates/aish-tools/src/{ => python}/python.rs | 26 +- crates/aish-tools/src/read_file/prompt.rs | 29 + crates/aish-tools/src/read_file/read_file.rs | 252 +++++ .../src/{ => secure_bash}/secure_bash.rs | 6 +- crates/aish-tools/src/skill_tool/prompt.rs | 25 + .../src/{ => skill_tool}/skill_tool.rs | 25 +- .../aish-tools/src/system_diagnose/prompt.rs | 21 + .../{ => system_diagnose}/system_diagnose.rs | 21 +- crates/aish-tools/src/web_fetch.rs | 935 ------------------ .../aish-tools/src/web_fetch/preapproved.rs | 127 +++ crates/aish-tools/src/web_fetch/prompt.rs | 34 + crates/aish-tools/src/web_fetch/utils.rs | 533 ++++++++++ crates/aish-tools/src/web_fetch/web_fetch.rs | 289 ++++++ crates/aish-tools/src/write_file/prompt.rs | 25 + .../aish-tools/src/write_file/write_file.rs | 106 ++ 52 files changed, 2871 insertions(+), 2544 deletions(-) rename crates/aish-tools/src/{ => ask_user}/ask_user.rs (82%) create mode 100644 crates/aish-tools/src/ask_user/prompt.rs rename crates/aish-tools/src/{ => bash}/bash.rs (97%) create mode 100644 crates/aish-tools/src/bash/prompt.rs rename crates/aish-tools/src/{ => channel_ask_user}/channel_ask_user.rs (71%) create mode 100644 crates/aish-tools/src/channel_ask_user/prompt.rs rename crates/aish-tools/src/{ => channel_bash}/channel_bash.rs (87%) create mode 100644 crates/aish-tools/src/channel_bash/prompt.rs create mode 100644 crates/aish-tools/src/edit_file/edit_file.rs create mode 100644 crates/aish-tools/src/edit_file/prompt.rs rename crates/aish-tools/src/{ => final_answer}/final_answer.rs (80%) create mode 100644 crates/aish-tools/src/final_answer/prompt.rs delete mode 100644 crates/aish-tools/src/fs.rs rename crates/aish-tools/src/{ => glob_tool}/glob_tool.rs (85%) create mode 100644 crates/aish-tools/src/glob_tool/prompt.rs rename crates/aish-tools/src/{ => grep_tool}/grep_tool.rs (87%) create mode 100644 crates/aish-tools/src/grep_tool/prompt.rs rename crates/aish-tools/src/{ => host_note}/host_note.rs (76%) create mode 100644 crates/aish-tools/src/host_note/prompt.rs rename crates/aish-tools/src/{ => memory_tool}/memory_tool.rs (79%) create mode 100644 crates/aish-tools/src/memory_tool/prompt.rs delete mode 100644 crates/aish-tools/src/plan_tool.rs create mode 100644 crates/aish-tools/src/plan_tool/enter_plan_mode.rs create mode 100644 crates/aish-tools/src/plan_tool/exit_plan_mode.rs create mode 100644 crates/aish-tools/src/plan_tool/list_plan_templates.rs create mode 100644 crates/aish-tools/src/plan_tool/prompt.rs create mode 100644 crates/aish-tools/src/python/prompt.rs rename crates/aish-tools/src/{ => python}/python.rs (88%) create mode 100644 crates/aish-tools/src/read_file/prompt.rs create mode 100644 crates/aish-tools/src/read_file/read_file.rs rename crates/aish-tools/src/{ => secure_bash}/secure_bash.rs (98%) create mode 100644 crates/aish-tools/src/skill_tool/prompt.rs rename crates/aish-tools/src/{ => skill_tool}/skill_tool.rs (86%) create mode 100644 crates/aish-tools/src/system_diagnose/prompt.rs rename crates/aish-tools/src/{ => system_diagnose}/system_diagnose.rs (92%) delete mode 100644 crates/aish-tools/src/web_fetch.rs create mode 100644 crates/aish-tools/src/web_fetch/preapproved.rs create mode 100644 crates/aish-tools/src/web_fetch/prompt.rs create mode 100644 crates/aish-tools/src/web_fetch/utils.rs create mode 100644 crates/aish-tools/src/web_fetch/web_fetch.rs create mode 100644 crates/aish-tools/src/write_file/prompt.rs create mode 100644 crates/aish-tools/src/write_file/write_file.rs diff --git a/crates/aish-i18n/locales/de-DE.yaml b/crates/aish-i18n/locales/de-DE.yaml index dce15e55..c207b2f1 100644 --- a/crates/aish-i18n/locales/de-DE.yaml +++ b/crates/aish-i18n/locales/de-DE.yaml @@ -353,25 +353,9 @@ security: tools: ask_user: - description: "Stellt dem Benutzer eine gezielte Rückfrage. Nur mit prompt wird Texteingabe verwendet; mit options werden Auswahlmöglichkeiten angeboten." unknown_kind: "Unbekannter Typ: {kind}" runtime_not_configured: "ask_user-Laufzeit ist nicht konfiguriert" execute_failed: "ask_user fehlgeschlagen: {error}" - param: - kind: "Interaktionstyp: text_input für freie Eingabe, choice_or_text für Auswahl mit eigener Eingabe" - prompt: "Die Frage, die dem Benutzer gestellt wird." - options: "Optionale Auswahlmöglichkeiten. Wenn sie fehlen oder leer sind, verwendet ask_user Texteingabe." - option_value: "Stabiler Optionswert in den Tool-Metadaten." - option_label: "Für den Benutzer sichtbare Bezeichnung in der Auswahlliste." - option_description: "Optionale Zusatzbeschreibung unterhalb der Option." - option_recommended: "Markiert diese Option sichtbar als empfohlene Vorgabe." - title: "Optionaler Titel für die Frage" - default: "Standardwert" - placeholder: "Platzhaltertext" - allow_freeform_input: "Erlaubt bei vorhandenen Optionen eine benutzerdefinierte Antwort. Standard ist true." - required: "Ob der Benutzer eine Antwort geben muss (Standard: true)" - allow_cancel: "Ob der Benutzer abbrechen/überspringen darf (Standard: true)" - min_length: "Minimale Länge für Texteingabe (Standard: 0)" validation: prompt_empty: "prompt darf nicht leer sein" option_value_empty: "option value darf nicht leer sein" @@ -388,10 +372,6 @@ tools: cancelled: "Benutzer hat die Frage abgebrochen." smart_log: - description: "Ermittelt automatisch Log-Quellen, analysiert Protokolle anhand der Anfrage und gibt eine passende Zusammenfassung zurück." - param: - query: "Ihre Diagnoseanfrage, z. B. 'ist nginx korrekt konfiguriert' oder 'gibt es verdächtige Anmeldungen'." - path: "Optional. Log-Pfad (Datei/Verzeichnis) oder ein systemd-Unit-Name. Wenn nicht angegeben, wird er aus der Anfrage abgeleitet." summary: file: "[DATEI] {ident}: {matched}/{total} Zeilen passend\nBeispiel:\n{sample}" dir: "[VERZ] {path}: {matched}/{total} Zeilen passend\nBeispiel:\n{sample}" diff --git a/crates/aish-i18n/locales/en-US.yaml b/crates/aish-i18n/locales/en-US.yaml index e3ea1645..51c6a034 100644 --- a/crates/aish-i18n/locales/en-US.yaml +++ b/crates/aish-i18n/locales/en-US.yaml @@ -632,10 +632,6 @@ security: tools: smart_log: - description: "Automatically determine log sources, analyze logs based on the query, and return a matched summary." - param: - query: "Your diagnostic query, e.g. 'is nginx configured correctly' or 'are there suspicious logins'." - path: "Optional. Log path (file/dir) or a systemd unit name. If omitted, it will be inferred from the query." summary: file: "[FILE] {ident}: matched {matched}/{total} lines\nSample:\n{sample}" dir: "[DIR] {path}: matched {matched}/{total} lines\nSample:\n{sample}" @@ -644,11 +640,7 @@ tools: no_logs_found: "No logs or matching information found." bash: - description: "Execute a bash command and return the output. Use this tool to run shell commands. IMPORTANT: Each command requires user confirmation. If the user rejects (N) or cancels (Ctrl+C) a command, you MUST NOT retry the same or similar command — acknowledge the cancellation and adjust your approach instead." security_handling_required: "bash command requires {level} security handling" - param: - command: "The bash command to execute" - timeout: "Timeout in seconds (default: 120)" missing_command: "Missing 'command' parameter" execute_failed: "Failed to execute: {error}" output_truncated: "[...{bytes} bytes truncated...]\n{tail}" @@ -656,11 +648,6 @@ tools: fs: read_file: - description: "Read the content of a file" - param: - path: "Path to the file to read" - offset: "Line offset to start reading from (0-based)" - limit: "Maximum number of lines to read" missing_path: "Missing 'path' parameter" read_failed: "Failed to read {path}: {error}" file_too_large: "File {path} is {size} bytes, exceeding the {limit} byte (32KB) limit" @@ -668,22 +655,12 @@ tools: empty_file: "(empty file)" offset_exceeds_length: "Offset {offset} exceeds file length ({length})" write_file: - description: "Write content to a file (creates or overwrites)" - param: - path: "Path to the file to write" - content: "Content to write" missing_path: "Missing 'path' parameter" missing_content: "Missing 'content' parameter" create_dirs_failed: "Failed to create parent dirs: {error}" write_success: "Wrote {bytes} bytes to {path}" write_failed: "Failed to write {path}: {error}" edit_file: - description: "Edit a file by replacing a specific string with a new string" - param: - path: "Path to the file" - old_string: "The text to replace" - new_string: "The replacement text" - replace_all: "Replace all occurrences (default: false)" missing_path: "Missing 'path' parameter" missing_old_string: "Missing 'old_string' parameter" missing_new_string: "Missing 'new_string' parameter" @@ -694,24 +671,8 @@ tools: edit_write_failed: "Failed to write {path}: {error}" ask_user: - description: "Ask the user a focused clarifying question. Use prompt for text input, and add options when you can offer choices." runtime_not_configured: "ask_user runtime is not configured" execute_failed: "ask_user failed: {error}" - param: - kind: "Interaction type: text_input for free-form, choice_or_text for options with custom input" - prompt: "The question to ask the user." - options: "Optional choices to offer the user. If omitted or empty, ask_user uses text input." - option_value: "Stable option value returned in tool metadata." - option_label: "User-facing label shown in the choice list." - option_description: "Optional extra detail shown below the option." - option_recommended: "Mark this option visibly as the recommended default." - title: "Optional title for the question" - default: "Default value" - placeholder: "Placeholder text" - allow_freeform_input: "When options are present, allow the user to choose Other and type a custom answer. Defaults to true." - required: "Whether the user must provide an answer (default: true)" - allow_cancel: "Whether the user can cancel/skip (default: true)" - min_length: "Minimum length for text input (default: 0)" validation: prompt_empty: "prompt cannot be empty" option_value_empty: "option value cannot be empty" @@ -742,13 +703,6 @@ tools: options_not_empty: "options must be a non-empty list for choice_or_text" memory: - description: "Search, store, or manage long-term memories. Use 'search' to find relevant past knowledge, 'store' to save important information, 'list' to see recent memories, 'forget' to remove outdated info." - param: - action: "Memory operation to perform" - query: "Search query (for 'search' action)" - content: "Content to store (for 'store' action)" - category: "Category for stored memory (default: other)" - memory_id: "Memory ID to forget (for 'forget' action)" missing_action: "Missing 'action' parameter" unknown_action: "Unknown action: {action}. Use search/store/forget/list." search_missing_query: "Missing 'query' for search" @@ -762,11 +716,6 @@ tools: not_available: "memory not available" host_note: - description: "Save, list, or delete notes about the current remote host. Use 'store' when the user tells you important facts about this server (services deployed, known issues, key paths, etc.). Use 'list' to review existing notes. Use 'forget' with a keyword to remove matching notes." - param: - action: "Note operation: store, list, or forget" - content: "Note content to save (for 'store' action)" - keyword: "Keyword to match notes for deletion (for 'forget' action)" missing_action: "Missing 'action' parameter" unknown_action: "Unknown action: {action}. Use store/list/forget." store_missing_content: "Missing 'content' for store" @@ -777,37 +726,20 @@ tools: forgot_none: "No matching notes found." python: - description: "Execute arbitrary Python code and return the result. Use print() for output." - param: - code: "The Python code to execute." missing_code: "Missing 'code' parameter" not_installed: "Python 3 is not installed or not in PATH." execute_failed: "Failed to execute Python: {error}" no_output: "Python code executed successfully with no output." grep: - description: "Search for patterns in files using regex" - param: - pattern: "Regular expression pattern to search for" - path: "File or directory to search in" - recursive: "Search recursively in directories" - case_insensitive: "Case insensitive search" - invert_match: "Show lines that do not match" - line_numbers: "Show line numbers" - count_only: "Only show count of matching lines" missing_pattern: "Missing 'pattern' parameter" invalid_regex: "Error: invalid regex pattern: {error}" glob: - description: "Find files matching glob patterns" - param: - pattern: "Glob pattern to match files" - path: "Base directory to search in (default: current directory)" missing_pattern: "Missing 'pattern' parameter" invalid_glob: "Error: invalid glob pattern: {error}" web_fetch: - description: "Fetch content from a specified URL, convert readable HTML to text, and answer a prompt about the page using a secondary model. IMPORTANT: WebFetch will fail for authenticated or private URLs. For GitHub URLs, prefer gh via bash when available." missing_url: "Missing 'url' parameter" missing_prompt: "Missing 'prompt' parameter" invalid_url: "Error: invalid URL. Provide a fully-qualified http or https URL without credentials." diff --git a/crates/aish-i18n/locales/es-ES.yaml b/crates/aish-i18n/locales/es-ES.yaml index 574db37d..d79acde7 100644 --- a/crates/aish-i18n/locales/es-ES.yaml +++ b/crates/aish-i18n/locales/es-ES.yaml @@ -353,25 +353,9 @@ security: tools: ask_user: - description: "Haz al usuario una pregunta de aclaración enfocada. Con solo prompt se usa texto libre; con options se ofrecen elecciones." unknown_kind: "Tipo desconocido: {kind}" runtime_not_configured: "el runtime de ask_user no está configurado" execute_failed: "ask_user falló: {error}" - param: - kind: "Tipo de interacción: text_input para texto libre, choice_or_text para opciones con entrada personalizada" - prompt: "La pregunta que se le hará al usuario." - options: "Opciones que se pueden ofrecer al usuario. Si faltan o están vacías, ask_user usa entrada de texto." - option_value: "Valor estable de la opción que se devuelve en los metadatos de la herramienta." - option_label: "Etiqueta visible para el usuario en la lista de opciones." - option_description: "Detalle opcional que se muestra debajo de la opción." - option_recommended: "Marca esta opción visiblemente como la recomendada." - title: "Título opcional para la pregunta" - default: "Valor predeterminado" - placeholder: "Texto de marcador" - allow_freeform_input: "Cuando hay opciones, permite elegir otra respuesta y escribirla. El valor predeterminado es true." - required: "Si el usuario debe proporcionar una respuesta (predeterminado: true)" - allow_cancel: "Si el usuario puede cancelar/omitir (predeterminado: true)" - min_length: "Longitud mínima para entrada de texto (predeterminado: 0)" validation: prompt_empty: "prompt no puede estar vacío" option_value_empty: "option value no puede estar vacío" @@ -388,10 +372,6 @@ tools: cancelled: "El usuario canceló la pregunta." smart_log: - description: "Determina automáticamente las fuentes de registros, analiza los logs según la consulta y devuelve un resumen coincidente." - param: - query: "Tu consulta de diagnóstico, por ejemplo 'nginx está bien configurado' o 'hay inicios de sesión sospechosos'." - path: "Opcional. Ruta del log (archivo/directorio) o nombre de una unidad systemd. Si se omite, se inferirá a partir de la consulta." summary: file: "[ARCHIVO] {ident}: {matched}/{total} líneas coincidentes\nMuestra:\n{sample}" dir: "[DIR] {path}: {matched}/{total} líneas coincidentes\nMuestra:\n{sample}" diff --git a/crates/aish-i18n/locales/fr-FR.yaml b/crates/aish-i18n/locales/fr-FR.yaml index 53ea5b88..41f2e0a7 100644 --- a/crates/aish-i18n/locales/fr-FR.yaml +++ b/crates/aish-i18n/locales/fr-FR.yaml @@ -353,25 +353,9 @@ security: tools: ask_user: - description: "Poser a l'utilisateur une question de clarification ciblee. Avec seulement prompt, ask_user utilise la saisie texte ; avec options, il propose des choix." unknown_kind: "Type inconnu : {kind}" runtime_not_configured: "le runtime ask_user n'est pas configure" execute_failed: "echec de ask_user : {error}" - param: - kind: "Type d'interaction : text_input pour saisie libre, choice_or_text pour options avec saisie personnalisee" - prompt: "La question a poser a l'utilisateur." - options: "Choix optionnels a proposer a l'utilisateur. S'ils sont absents ou vides, ask_user utilise la saisie texte." - option_value: "Valeur stable de l'option retournee dans les metadonnees de l'outil." - option_label: "Libelle affiche a l'utilisateur dans la liste." - option_description: "Detail optionnel affiche sous l'option." - option_recommended: "Marque visiblement cette option comme recommandee." - title: "Titre optionnel pour la question" - default: "Valeur par defaut" - placeholder: "Texte indicatif" - allow_freeform_input: "Quand des options sont presentes, autorise une reponse personnalisee. La valeur par defaut est true." - required: "Indique si l'utilisateur doit fournir une reponse (par defaut : true)" - allow_cancel: "Indique si l'utilisateur peut annuler/ignorer (par defaut : true)" - min_length: "Longueur minimale pour la saisie texte (par defaut : 0)" validation: prompt_empty: "prompt ne peut pas etre vide" option_value_empty: "option value ne peut pas etre vide" @@ -388,10 +372,6 @@ tools: cancelled: "L'utilisateur a annule la question." smart_log: - description: "Determine automatiquement les sources de journaux, analyse les logs selon la requete et renvoie un resume correspondant." - param: - query: "Votre requete de diagnostic, par exemple 'nginx est-il correctement configure' ou 'y a-t-il des connexions suspectes'." - path: "Optionnel. Chemin du journal (fichier/repertoire) ou nom d'une unite systemd. S'il est omis, il sera deduit de la requete." summary: file: "[FICHIER] {ident} : {matched}/{total} lignes correspondantes\nExemple :\n{sample}" dir: "[REP] {path} : {matched}/{total} lignes correspondantes\nExemple :\n{sample}" diff --git a/crates/aish-i18n/locales/ja-JP.yaml b/crates/aish-i18n/locales/ja-JP.yaml index df6ee47e..f6010d59 100644 --- a/crates/aish-i18n/locales/ja-JP.yaml +++ b/crates/aish-i18n/locales/ja-JP.yaml @@ -353,25 +353,9 @@ security: tools: ask_user: - description: "ユーザーに焦点を絞った確認質問をします。prompt のみならテキスト入力、options があれば選択肢を提示します。" unknown_kind: "不明な種類です: {kind}" runtime_not_configured: "ask_user ランタイムが構成されていません" execute_failed: "ask_user に失敗しました: {error}" - param: - kind: "操作タイプ: text_input は自由入力、choice_or_text は選択肢とカスタム入力" - prompt: "ユーザーに尋ねる質問です。" - options: "ユーザーに提示する任意の選択肢です。省略または空の場合、ask_user はテキスト入力を使います。" - option_value: "ツールメタデータに返される安定した選択肢の値です。" - option_label: "選択リストに表示されるユーザー向けラベルです。" - option_description: "選択肢の下に表示される任意の補足説明です。" - option_recommended: "この選択肢を推奨として明示表示します。" - title: "質問の任意タイトル" - default: "既定値" - placeholder: "プレースホルダーテキスト" - allow_freeform_input: "選択肢がある場合にカスタム回答を許可します。既定値は true です。" - required: "ユーザーが回答を入力する必要があるかどうか(既定値: true)" - allow_cancel: "ユーザーがキャンセル/スキップできるかどうか(既定値: true)" - min_length: "テキスト入力の最小文字数(既定値: 0)" validation: prompt_empty: "prompt は空にできません" option_value_empty: "option value は空にできません" @@ -388,10 +372,6 @@ tools: cancelled: "ユーザーが質問をキャンセルしました。" smart_log: - description: "ログソースを自動判定し、問い合わせに基づいてログを解析し、一致する要約を返します。" - param: - query: "診断したい内容。例: nginx は正しく設定されているか、怪しいログインはあるか。" - path: "任意。ログパス(ファイルまたはディレクトリ)または systemd ユニット名。省略時は問い合わせから推定します。" summary: file: "[FILE] {ident}: 一致 {matched}/{total} 行\nサンプル:\n{sample}" dir: "[DIR] {path}: 一致 {matched}/{total} 行\nサンプル:\n{sample}" diff --git a/crates/aish-i18n/locales/zh-CN.yaml b/crates/aish-i18n/locales/zh-CN.yaml index c79e25a8..3e7bd484 100644 --- a/crates/aish-i18n/locales/zh-CN.yaml +++ b/crates/aish-i18n/locales/zh-CN.yaml @@ -631,10 +631,6 @@ security: tools: smart_log: - description: "自动判断日志来源,根据查询分析日志,并返回匹配摘要。" - param: - query: "你的诊断问题,例如\"nginx 配置是否正确\"或\"是否存在可疑登录\"。" - path: "可选。日志路径(文件/目录)或 systemd 单元名。若为空,将根据查询自动推断。" summary: file: "[文件] {ident}:匹配 {matched}/{total} 行\n示例:\n{sample}" dir: "[目录] {path}:匹配 {matched}/{total} 行\n示例:\n{sample}" @@ -643,11 +639,7 @@ tools: no_logs_found: "未找到任何日志或匹配信息。" bash: - description: "执行 bash 命令并返回输出。使用此工具运行 shell 命令。重要:每条命令都需要用户确认。如果用户拒绝(N)或取消(Ctrl+C),你绝对不能重试相同或类似的命令——应确认取消并调整方案。" security_handling_required: "bash 命令需要 {level} 级安全处理" - param: - command: "要执行的 bash 命令" - timeout: "超时时间(秒)(默认:120)" missing_command: "缺少 'command' 参数" execute_failed: "执行失败: {error}" output_truncated: "[...截断了 {bytes} 字节...]\n{tail}" @@ -655,11 +647,6 @@ tools: fs: read_file: - description: "读取文件内容" - param: - path: "要读取的文件路径" - offset: "开始读取的行偏移(从 0 开始)" - limit: "要读取的最大行数" missing_path: "缺少 'path' 参数" read_failed: "读取 {path} 失败: {error}" file_too_large: "文件 {path} 为 {size} 字节,超过了 {limit} 字节(32KB)的限制" @@ -667,22 +654,12 @@ tools: empty_file: "(空文件)" offset_exceeds_length: "偏移量 {offset} 超过了文件长度({length})" write_file: - description: "将内容写入文件(创建或覆盖)" - param: - path: "要写入的文件路径" - content: "要写入的内容" missing_path: "缺少 'path' 参数" missing_content: "缺少 'content' 参数" create_dirs_failed: "创建父目录失败: {error}" write_success: "已写入 {bytes} 字节到 {path}" write_failed: "写入 {path} 失败: {error}" edit_file: - description: "通过替换特定字符串来编辑文件" - param: - path: "文件路径" - old_string: "要替换的文本" - new_string: "替换后的文本" - replace_all: "替换所有出现的位置(默认:false)" missing_path: "缺少 'path' 参数" missing_old_string: "缺少 'old_string' 参数" missing_new_string: "缺少 'new_string' 参数" @@ -693,24 +670,8 @@ tools: edit_write_failed: "写入 {path} 失败: {error}" ask_user: - description: "向用户提出一个聚焦的澄清问题。只填 prompt 时为文本输入,提供 options 时为选项交互。" runtime_not_configured: "ask_user runtime 未配置" execute_failed: "ask_user 执行失败:{error}" - param: - kind: "交互类型:text_input 用于自由输入,choice_or_text 用于选项加自定义输入" - prompt: "要向用户提出的问题。" - options: "可选的候选项。为空或省略时,ask_user 使用文本输入。" - option_value: "返回到工具元数据中的稳定选项值。" - option_label: "展示给用户的选项标签。" - option_description: "显示在选项下方的附加说明。" - option_recommended: "显式标记该选项为推荐默认项。" - title: "问题的可选标题" - default: "默认值" - placeholder: "占位符文本" - allow_freeform_input: "当存在 options 时,允许用户选择“其他”并输入自定义答案,默认 true。" - required: "用户是否必须提供答案(默认:true)" - allow_cancel: "用户是否可以取消/跳过(默认:true)" - min_length: "文本输入的最小长度(默认:0)" validation: prompt_empty: "prompt 不能为空" option_value_empty: "option value 不能为空" @@ -741,13 +702,6 @@ tools: options_not_empty: "choice_or_text 的 options 必须是非空列表" memory: - description: "搜索、存储或管理长期记忆。使用 'search' 查找相关的过往知识,'store' 保存重要信息,'list' 查看最近的记忆,'forget' 删除过时信息。" - param: - action: "要执行的记忆操作" - query: "搜索查询(用于 'search' 操作)" - content: "要存储的内容(用于 'store' 操作)" - category: "存储记忆的类别(默认:other)" - memory_id: "要忘记的记忆 ID(用于 'forget' 操作)" missing_action: "缺少 'action' 参数" unknown_action: "未知操作: {action}。请使用 search/store/forget/list。" search_missing_query: "search 缺少 'query' 参数" @@ -761,11 +715,6 @@ tools: not_available: "记忆功能不可用" host_note: - description: "保存、查看或删除当前远程主机的备注。当用户告诉你关于此服务器的重要信息(部署的服务、已知问题、关键路径等)时使用 'store' 保存。使用 'list' 查看已有备注。使用 'forget' 并提供关键词删除匹配的备注。" - param: - action: "备注操作: store、list 或 forget" - content: "要保存的备注内容(用于 'store' 操作)" - keyword: "用于匹配删除备注的关键词(用于 'forget' 操作)" missing_action: "缺少 'action' 参数" unknown_action: "未知操作: {action}。请使用 store/list/forget。" store_missing_content: "store 缺少 'content' 参数" @@ -776,37 +725,20 @@ tools: forgot_none: "未找到匹配的备注。" python: - description: "执行任意 Python 代码并返回结果。使用 print() 输出。" - param: - code: "要执行的 Python 代码。" missing_code: "缺少 'code' 参数" not_installed: "未安装 Python 3 或不在 PATH 中。" execute_failed: "执行 Python 失败: {error}" no_output: "Python 代码执行成功,无输出。" grep: - description: "使用正则表达式在文件中搜索模式" - param: - pattern: "要搜索的正则表达式模式" - path: "要搜索的文件或目录" - recursive: "在目录中递归搜索" - case_insensitive: "不区分大小写搜索" - invert_match: "显示不匹配的行" - line_numbers: "显示行号" - count_only: "仅显示匹配行的计数" missing_pattern: "缺少 'pattern' 参数" invalid_regex: "错误:无效的正则表达式: {error}" glob: - description: "查找匹配 glob 模式的文件" - param: - pattern: "匹配文件的 glob 模式" - path: "搜索的基础目录(默认:当前目录)" missing_pattern: "缺少 'pattern' 参数" invalid_glob: "错误:无效的 glob 模式: {error}" web_fetch: - description: "从指定 URL 抓取内容,将可读 HTML 转为文本,并使用二级模型按 prompt 分析页面。重要:WebFetch 无法访问需要认证或私有的 URL。对于 GitHub URL,如可用请优先通过 bash 使用 gh。" missing_url: "缺少 'url' 参数" missing_prompt: "缺少 'prompt' 参数" invalid_url: "错误:无效 URL。请提供完整的 http 或 https URL,且不要包含凭据。" diff --git a/crates/aish-llm/src/agent.rs b/crates/aish-llm/src/agent.rs index 25a4647e..10701dc0 100644 --- a/crates/aish-llm/src/agent.rs +++ b/crates/aish-llm/src/agent.rs @@ -255,7 +255,8 @@ impl<'a> ReActAgent<'a> { cancel.reset(); let mut messages: Vec = Vec::new(); - messages.push(ChatMessage::system(system_prompt)); + let system_prompt = self.session.system_prompt_with_tool_prompts(system_prompt); + messages.push(ChatMessage::system(&system_prompt)); messages.push(ChatMessage::user(query)); let tool_specs: Vec = self.session.tool_specs(); diff --git a/crates/aish-llm/src/session.rs b/crates/aish-llm/src/session.rs index 34a730ee..5a707fe1 100644 --- a/crates/aish-llm/src/session.rs +++ b/crates/aish-llm/src/session.rs @@ -164,6 +164,53 @@ impl LlmSession { } } + fn tool_visible_in_phase(tool_name: &str, phase: &PlanPhase) -> bool { + match phase { + PlanPhase::Normal => true, + PlanPhase::Planning => aish_core::PLANNING_VISIBLE_TOOLS.contains(&tool_name), + } + } + + pub fn filtered_tool_prompt_section(&self) -> Option { + let phase = self.plan_state.lock().unwrap().phase.clone(); + let mut prompts: Vec<(&str, &str)> = self + .tools + .values() + .filter(|tool| Self::tool_visible_in_phase(tool.name(), &phase)) + .filter_map(|tool| { + let prompt = tool.prompt().trim(); + if prompt.is_empty() { + None + } else { + Some((tool.name(), prompt)) + } + }) + .collect(); + + if prompts.is_empty() { + return None; + } + + prompts.sort_by(|a, b| a.0.cmp(b.0)); + + let mut section = String::from("## Tool Instructions\n"); + for (name, prompt) in prompts { + section.push_str("\n### "); + section.push_str(name); + section.push('\n'); + section.push_str(prompt); + section.push('\n'); + } + Some(section) + } + + pub fn system_prompt_with_tool_prompts(&self, system_prompt: &str) -> String { + match self.filtered_tool_prompt_section() { + Some(section) => format!("{}\n\n{}", system_prompt.trim_end(), section), + None => system_prompt.to_string(), + } + } + /// Get a reference to the plan state (for external coordination). pub fn plan_state(&self) -> Arc> { Arc::clone(&self.plan_state) @@ -286,7 +333,9 @@ impl LlmSession { // Build initial message list let mut messages: Vec = Vec::new(); if let Some(sys) = system_message { - messages.push(ChatMessage::system(sys)); + messages.push(ChatMessage::system( + &self.system_prompt_with_tool_prompts(sys), + )); } messages.extend_from_slice(context_messages); messages.push(user_msg.clone()); @@ -2499,12 +2548,21 @@ mod tests { // Mock tool for testing struct MockTool { name: String, + prompt: String, } impl MockTool { fn new(name: &str) -> Self { Self { name: name.to_string(), + prompt: String::new(), + } + } + + fn with_prompt(name: &str, prompt: &str) -> Self { + Self { + name: name.to_string(), + prompt: prompt.to_string(), } } } @@ -2522,11 +2580,70 @@ mod tests { serde_json::json!({"type": "object", "properties": {}}) } + fn prompt(&self) -> &str { + &self.prompt + } + fn execute(&self, _args: serde_json::Value) -> crate::types::ToolResult { crate::types::ToolResult::success("mock result") } } + #[test] + fn test_tool_prompt_section_uses_only_non_empty_prompts() { + let mut session = LlmSession::new("http://localhost", "key", "model", None, None); + session.register_tool(Box::new(MockTool::new("empty_tool"))); + session.register_tool(Box::new(MockTool::with_prompt( + "prompt_tool", + "Use carefully.", + ))); + + let section = session.filtered_tool_prompt_section().unwrap(); + + assert!(section.contains("## Tool Instructions")); + assert!(section.contains("### prompt_tool")); + assert!(section.contains("Use carefully.")); + assert!(!section.contains("empty_tool")); + } + + #[test] + fn test_tool_prompt_section_respects_planning_filter() { + let mut session = LlmSession::new("http://localhost", "key", "model", None, None); + session.register_tool(Box::new(MockTool::with_prompt( + "read_file", + "Read files during planning.", + ))); + session.register_tool(Box::new(MockTool::with_prompt( + "bash_exec", + "Run commands.", + ))); + + { + let mut state = session.plan_state.lock().unwrap(); + state.phase = aish_core::PlanPhase::Planning; + } + + let section = session.filtered_tool_prompt_section().unwrap(); + + assert!(section.contains("### read_file")); + assert!(section.contains("Read files during planning.")); + assert!(!section.contains("bash_exec")); + assert!(!section.contains("Run commands.")); + } + + #[test] + fn test_system_prompt_with_tool_prompts_appends_section() { + let mut session = LlmSession::new("http://localhost", "key", "model", None, None); + session.register_tool(Box::new(MockTool::with_prompt("mock", "Mock guidance."))); + + let system_prompt = session.system_prompt_with_tool_prompts("Base prompt.\n"); + + assert!(system_prompt.starts_with("Base prompt.")); + assert!(system_prompt.contains("## Tool Instructions")); + assert!(system_prompt.contains("### mock")); + assert!(system_prompt.contains("Mock guidance.")); + } + #[test] fn test_tool_filtering_with_registered_tools() { use aish_core::PlanPhase; diff --git a/crates/aish-llm/src/types.rs b/crates/aish-llm/src/types.rs index 8f5eefd7..9b11ecbc 100644 --- a/crates/aish-llm/src/types.rs +++ b/crates/aish-llm/src/types.rs @@ -505,6 +505,10 @@ pub trait Tool: Send + Sync { fn description(&self) -> &str; fn parameters(&self) -> serde_json::Value; + fn prompt(&self) -> &str { + "" + } + fn to_spec(&self) -> ToolSpec { ToolSpec { r#type: "function".into(), diff --git a/crates/aish-tools/src/ask_user.rs b/crates/aish-tools/src/ask_user/ask_user.rs similarity index 82% rename from crates/aish-tools/src/ask_user.rs rename to crates/aish-tools/src/ask_user/ask_user.rs index df77d2a0..3bcb54a9 100644 --- a/crates/aish-tools/src/ask_user.rs +++ b/crates/aish-tools/src/ask_user/ask_user.rs @@ -1,10 +1,12 @@ use std::io; -use std::sync::{Arc, OnceLock}; +use std::sync::Arc; use aish_i18n::{t, t_with_args}; use aish_llm::{Tool, ToolResult}; use serde::Deserialize; +use super::prompt; + #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] pub struct AskUserOption { pub value: String, @@ -77,82 +79,12 @@ fn default_allow_freeform_input() -> bool { true } -static DESCRIPTION: OnceLock = OnceLock::new(); - pub(crate) fn ask_user_description() -> &'static str { - DESCRIPTION.get_or_init(|| t("tools.ask_user.description")) + prompt::DESCRIPTION } pub(crate) fn ask_user_parameters() -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "prompt": { - "type": "string", - "description": t("tools.ask_user.param.prompt") - }, - "options": { - "type": "array", - "description": t("tools.ask_user.param.options"), - "items": { - "type": "object", - "properties": { - "value": { - "type": "string", - "description": t("tools.ask_user.param.option_value") - }, - "label": { - "type": "string", - "description": t("tools.ask_user.param.option_label") - }, - "description": { - "type": "string", - "description": t("tools.ask_user.param.option_description") - }, - "recommended": { - "type": "boolean", - "description": t("tools.ask_user.param.option_recommended") - } - }, - "required": ["value", "label"] - } - }, - "title": { - "type": "string", - "description": t("tools.ask_user.param.title") - }, - "default": { - "type": "string", - "description": t("tools.ask_user.param.default") - }, - "placeholder": { - "type": "string", - "description": t("tools.ask_user.param.placeholder") - }, - "allow_freeform_input": { - "type": "boolean", - "description": t("tools.ask_user.param.allow_freeform_input"), - "default": true - }, - "required": { - "type": "boolean", - "description": t("tools.ask_user.param.required"), - "default": true - }, - "allow_cancel": { - "type": "boolean", - "description": t("tools.ask_user.param.allow_cancel"), - "default": true - }, - "min_length": { - "type": "integer", - "minimum": 0, - "description": t("tools.ask_user.param.min_length"), - "default": 0 - } - }, - "required": ["prompt"] - }) + prompt::parameters() } pub struct AskUserTool { @@ -194,6 +126,10 @@ impl Tool for AskUserTool { ask_user_parameters() } + fn prompt(&self) -> &str { + prompt::PROMPT + } + fn execute(&self, args: serde_json::Value) -> ToolResult { let request = match parse_args(args) { Ok(request) => request, diff --git a/crates/aish-tools/src/ask_user/prompt.rs b/crates/aish-tools/src/ask_user/prompt.rs new file mode 100644 index 00000000..49b16ac7 --- /dev/null +++ b/crates/aish-tools/src/ask_user/prompt.rs @@ -0,0 +1,81 @@ +pub(crate) const DESCRIPTION: &str = "Ask the user a focused clarifying question."; + +pub(crate) const PROMPT: &str = r#"Use this tool only when a small amount of user input is needed to continue. + +Usage: +- Ask one focused question at a time. +- Prefer options when the likely answers are known. +- Use allow_freeform_input=false only when the user must choose from the provided options. +- Do not ask for secrets such as passwords, API keys, or tokens."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "The question to ask the user." + }, + "options": { + "type": "array", + "description": "Optional choices to offer the user. If omitted or empty, ask_user uses text input.", + "items": { + "type": "object", + "properties": { + "value": { + "type": "string", + "description": "Stable option value returned in tool metadata." + }, + "label": { + "type": "string", + "description": "User-facing label shown in the choice list." + }, + "description": { + "type": "string", + "description": "Optional extra detail shown below the option." + }, + "recommended": { + "type": "boolean", + "description": "Mark this option visibly as the recommended default." + } + }, + "required": ["value", "label"] + } + }, + "title": { + "type": "string", + "description": "Optional title for the question." + }, + "default": { + "type": "string", + "description": "Default value." + }, + "placeholder": { + "type": "string", + "description": "Placeholder text." + }, + "allow_freeform_input": { + "type": "boolean", + "description": "When options are present, allow the user to choose Other and type a custom answer. Defaults to true.", + "default": true + }, + "required": { + "type": "boolean", + "description": "Whether the user must provide an answer. Defaults to true.", + "default": true + }, + "allow_cancel": { + "type": "boolean", + "description": "Whether the user can cancel or skip the question. Defaults to true.", + "default": true + }, + "min_length": { + "type": "integer", + "minimum": 0, + "description": "Minimum length for text input. Defaults to 0.", + "default": 0 + } + }, + "required": ["prompt"] + }) +} diff --git a/crates/aish-tools/src/bash.rs b/crates/aish-tools/src/bash/bash.rs similarity index 97% rename from crates/aish-tools/src/bash.rs rename to crates/aish-tools/src/bash/bash.rs index a1a89eeb..ebb3926e 100644 --- a/crates/aish-tools/src/bash.rs +++ b/crates/aish-tools/src/bash/bash.rs @@ -13,6 +13,8 @@ use aish_security::{ load_policy, secret::SecretVault, SecurityDecision, SecurityManager, SecurityRequest, }; +use super::prompt; + /// Large keep_bytes for the silent PTY executor to capture full command output. /// The BashOutputOffload will handle threshold-based truncation and disk offload. const CAPTURE_KEEP_BYTES: usize = 10 * 1024 * 1024; // 10MB @@ -108,13 +110,6 @@ pub struct BashTool { secret_vault: VaultSlot, } -/// Cached translated description. -static DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.bash.description")) -} - fn timeout_secs(args: &serde_json::Value) -> Result, ToolResult> { match args.get("timeout") { None => Ok(None), @@ -483,25 +478,15 @@ impl Tool for BashTool { } fn description(&self) -> &str { - get_description() + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "command": { - "type": "string", - "description": "The bash command to execute" - }, - "timeout": { - "type": "integer", - "minimum": 1, - "description": "Timeout in seconds. If omitted, the command runs until completion or cancellation." - } - }, - "required": ["command"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn preflight(&self, args: &serde_json::Value) -> PreflightResult { diff --git a/crates/aish-tools/src/bash/prompt.rs b/crates/aish-tools/src/bash/prompt.rs new file mode 100644 index 00000000..51df992c --- /dev/null +++ b/crates/aish-tools/src/bash/prompt.rs @@ -0,0 +1,27 @@ +pub(crate) const DESCRIPTION: &str = "Execute a bash command and return the output."; + +pub(crate) const PROMPT: &str = r#"Use this tool to run shell commands. + +Usage: +- Explain non-trivial commands before running them. +- Prefer read_file, grep, or glob when those tools directly match the task. +- Use timeout only when a bounded runtime is expected. +- Do not retry commands the user rejected or cancelled."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "Bash command to execute." + }, + "timeout": { + "type": "integer", + "minimum": 1, + "description": "Timeout in seconds. If omitted, the command runs until completion or cancellation." + } + }, + "required": ["command"] + }) +} diff --git a/crates/aish-tools/src/channel_ask_user.rs b/crates/aish-tools/src/channel_ask_user/channel_ask_user.rs similarity index 71% rename from crates/aish-tools/src/channel_ask_user.rs rename to crates/aish-tools/src/channel_ask_user/channel_ask_user.rs index e7385b95..dc91dd96 100644 --- a/crates/aish-tools/src/channel_ask_user.rs +++ b/crates/aish-tools/src/channel_ask_user/channel_ask_user.rs @@ -8,12 +8,7 @@ use aish_llm::{Tool, ToolResult}; use aish_pty::{AiEvent, AskUserAnswer, AskUserOption, AskUserRequest}; -/// Shared translated description — same as AskUserTool. -static DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.ask_user.description")) -} +use super::prompt; pub struct ChannelAskUserTool { question_sender: std::sync::mpsc::Sender, @@ -65,56 +60,15 @@ impl Tool for ChannelAskUserTool { } fn description(&self) -> &str { - get_description() + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "kind": { - "type": "string", - "enum": ["text_input", "choice_or_text"], - "description": "Interaction type: text_input for free-form, choice_or_text for options with custom input" - }, - "prompt": { - "type": "string", - "description": "The question to ask the user" - }, - "options": { - "type": "array", - "description": "Predefined options for choice_or_text", - "items": { - "type": "object", - "properties": { - "value": {"type": "string"}, - "label": {"type": "string"}, - "description": {"type": "string"} - }, - "required": ["value", "label"] - } - }, - "title": { - "type": "string", - "description": "Optional title for the question" - }, - "default": { - "type": "string", - "description": "Default value" - }, - "allow_cancel": { - "type": "boolean", - "description": "Whether the user can cancel/skip (default: true)", - "default": true - }, - "min_length": { - "type": "integer", - "description": "Minimum length for text input (default: 0)", - "default": 0 - } - }, - "required": ["kind", "prompt"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/channel_ask_user/prompt.rs b/crates/aish-tools/src/channel_ask_user/prompt.rs new file mode 100644 index 00000000..5fccee7a --- /dev/null +++ b/crates/aish-tools/src/channel_ask_user/prompt.rs @@ -0,0 +1,57 @@ +pub(crate) const DESCRIPTION: &str = "Ask the user a focused clarifying question."; + +pub(crate) const PROMPT: &str = r#"Use this tool only when a small amount of user input is needed to continue. + +Usage: +- Ask one focused question at a time. +- Prefer options when the likely answers are known. +- Do not ask for secrets such as passwords, API keys, or tokens."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "kind": { + "type": "string", + "enum": ["text_input", "choice_or_text"], + "description": "Interaction type: text_input for free-form, choice_or_text for options with custom input." + }, + "prompt": { + "type": "string", + "description": "Question to ask the user." + }, + "options": { + "type": "array", + "description": "Predefined options for choice_or_text.", + "items": { + "type": "object", + "properties": { + "value": {"type": "string"}, + "label": {"type": "string"}, + "description": {"type": "string"} + }, + "required": ["value", "label"] + } + }, + "title": { + "type": "string", + "description": "Optional title for the question." + }, + "default": { + "type": "string", + "description": "Default value." + }, + "allow_cancel": { + "type": "boolean", + "description": "Whether the user can cancel or skip. Defaults to true.", + "default": true + }, + "min_length": { + "type": "integer", + "description": "Minimum length for text input. Defaults to 0.", + "default": 0 + } + }, + "required": ["kind", "prompt"] + }) +} diff --git a/crates/aish-tools/src/channel_bash.rs b/crates/aish-tools/src/channel_bash/channel_bash.rs similarity index 87% rename from crates/aish-tools/src/channel_bash.rs rename to crates/aish-tools/src/channel_bash/channel_bash.rs index 3b1db199..d2313d89 100644 --- a/crates/aish-tools/src/channel_bash.rs +++ b/crates/aish-tools/src/channel_bash/channel_bash.rs @@ -10,11 +10,7 @@ use aish_llm::{Tool, ToolResult}; use aish_pty::{truncate_utf8_safe, AiEvent, BashExecResult}; -static DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.bash.description")) -} +use super::prompt; pub struct ChannelBashTool { event_sender: std::sync::mpsc::Sender, @@ -32,25 +28,15 @@ impl Tool for ChannelBashTool { } fn description(&self) -> &str { - get_description() + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "command": { - "type": "string", - "description": aish_i18n::t("tools.bash.param.command") - }, - "timeout": { - "type": "integer", - "description": aish_i18n::t("tools.bash.param.timeout"), - "default": 120 - } - }, - "required": ["command"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/channel_bash/prompt.rs b/crates/aish-tools/src/channel_bash/prompt.rs new file mode 100644 index 00000000..7b2231d1 --- /dev/null +++ b/crates/aish-tools/src/channel_bash/prompt.rs @@ -0,0 +1,26 @@ +pub(crate) const DESCRIPTION: &str = "Execute a bash command and return the output."; + +pub(crate) const PROMPT: &str = r#"Use this tool to run shell commands in the current SSH-backed session. + +Usage: +- Explain non-trivial commands before running them. +- Use timeout only when a bounded runtime is expected. +- Do not retry commands the user rejected or cancelled."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "Bash command to execute." + }, + "timeout": { + "type": "integer", + "description": "Timeout in seconds.", + "default": 120 + } + }, + "required": ["command"] + }) +} diff --git a/crates/aish-tools/src/edit_file/edit_file.rs b/crates/aish-tools/src/edit_file/edit_file.rs new file mode 100644 index 00000000..f082bb54 --- /dev/null +++ b/crates/aish-tools/src/edit_file/edit_file.rs @@ -0,0 +1,173 @@ +use aish_i18n; +use aish_llm::{Tool, ToolResult}; + +use super::prompt; + +/// Edit file tool (string replacement). +pub struct EditFileTool; + +impl Default for EditFileTool { + fn default() -> Self { + Self::new() + } +} + +impl EditFileTool { + pub fn new() -> Self { + Self + } +} + +impl Tool for EditFileTool { + fn name(&self) -> &str { + "edit_file" + } + + fn description(&self) -> &str { + prompt::DESCRIPTION + } + + fn parameters(&self) -> serde_json::Value { + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT + } + + fn execute(&self, args: serde_json::Value) -> ToolResult { + let path = match args.get("path").and_then(|p| p.as_str()) { + Some(p) => p, + None => return ToolResult::error(aish_i18n::t("tools.fs.edit_file.missing_path")), + }; + let old = match args.get("old_string").and_then(|o| o.as_str()) { + Some(o) => o, + None => { + return ToolResult::error(aish_i18n::t("tools.fs.edit_file.missing_old_string")) + } + }; + let new = match args.get("new_string").and_then(|n| n.as_str()) { + Some(n) => n, + None => { + return ToolResult::error(aish_i18n::t("tools.fs.edit_file.missing_new_string")) + } + }; + let replace_all = args + .get("replace_all") + .and_then(|r| r.as_bool()) + .unwrap_or(false); + + let content = match std::fs::read_to_string(path) { + Ok(c) => c, + Err(e) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("error".to_string(), e.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.edit_file.edit_read_failed", + &args_map, + )); + } + }; + + if !content.contains(old) { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.edit_file.old_string_not_found", + &args_map, + )); + } + + let new_content = if replace_all { + content.replace(old, new) + } else { + let count = content.matches(old).count(); + if count > 1 { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("count".to_string(), count.to_string()); + args_map.insert("path".to_string(), path.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.edit_file.old_string_ambiguous", + &args_map, + )); + } + content.replacen(old, new, 1) + }; + + match std::fs::write(path, new_content) { + Ok(()) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + ToolResult::success(aish_i18n::t_with_args( + "tools.fs.edit_file.edit_success", + &args_map, + )) + } + Err(e) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("error".to_string(), e.to_string()); + ToolResult::error(aish_i18n::t_with_args( + "tools.fs.edit_file.edit_write_failed", + &args_map, + )) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aish_llm::Tool; + use std::fs; + + fn temp_dir() -> tempfile::TempDir { + tempfile::tempdir().expect("failed to create temp dir") + } + + #[test] + fn test_edit_file_replace_all() { + let dir = temp_dir(); + let file_path = dir.path().join("test.txt"); + fs::write(&file_path, "foo bar foo baz foo").unwrap(); + + let tool = EditFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap(), + "old_string": "foo", + "new_string": "qux", + "replace_all": true + })); + + assert!(result.ok); + let content = fs::read_to_string(&file_path).unwrap(); + assert_eq!(content, "qux bar qux baz qux"); + } + + #[test] + fn test_edit_file_uniqueness_check() { + aish_i18n::set_locale("en-US"); + + let dir = temp_dir(); + let file_path = dir.path().join("test.txt"); + fs::write(&file_path, "foo bar foo baz").unwrap(); + + let tool = EditFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap(), + "old_string": "foo", + "new_string": "qux" + })); + + assert!(!result.ok); + assert!( + result.output.contains("times") || result.output.contains("ambiguous"), + "Expected uniqueness error, got: {}", + result.output + ); + let content = fs::read_to_string(&file_path).unwrap(); + assert_eq!(content, "foo bar foo baz"); + } +} diff --git a/crates/aish-tools/src/edit_file/prompt.rs b/crates/aish-tools/src/edit_file/prompt.rs new file mode 100644 index 00000000..41b7c774 --- /dev/null +++ b/crates/aish-tools/src/edit_file/prompt.rs @@ -0,0 +1,33 @@ +pub(crate) const DESCRIPTION: &str = "Edit a file by replacing exact text."; + +pub(crate) const PROMPT: &str = r#"Use this tool to make exact string replacements in text files. + +Usage: +- old_string must match exactly. +- Provide enough surrounding context when replacing a repeated string. +- Use replace_all only when every occurrence should change."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file." + }, + "old_string": { + "type": "string", + "description": "Exact text to replace." + }, + "new_string": { + "type": "string", + "description": "Replacement text." + }, + "replace_all": { + "type": "boolean", + "description": "Replace all occurrences. Defaults to false." + } + }, + "required": ["path", "old_string", "new_string"] + }) +} diff --git a/crates/aish-tools/src/final_answer.rs b/crates/aish-tools/src/final_answer/final_answer.rs similarity index 80% rename from crates/aish-tools/src/final_answer.rs rename to crates/aish-tools/src/final_answer/final_answer.rs index 0e5e497f..3bb4dbb5 100644 --- a/crates/aish-tools/src/final_answer.rs +++ b/crates/aish-tools/src/final_answer/final_answer.rs @@ -1,5 +1,7 @@ use aish_llm::{Tool, ToolResult}; +use super::prompt; + /// Tool that signals the agent has reached a final answer. /// /// When the LLM calls this tool, the agent loop should terminate and return @@ -25,20 +27,15 @@ impl Tool for FinalAnswerTool { } fn description(&self) -> &str { - "Submit the final answer to the user's question. Call this tool when you have completed your analysis and have a definitive answer. The answer will be shown to the user directly." + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "answer": { - "type": "string", - "description": "The complete final answer to present to the user" - } - }, - "required": ["answer"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/final_answer/prompt.rs b/crates/aish-tools/src/final_answer/prompt.rs new file mode 100644 index 00000000..8827db96 --- /dev/null +++ b/crates/aish-tools/src/final_answer/prompt.rs @@ -0,0 +1,20 @@ +pub(crate) const DESCRIPTION: &str = "Submit the final answer to the user's question."; + +pub(crate) const PROMPT: &str = r#"Use this tool when the task is complete and the final answer is ready. + +Usage: +- Put the complete user-facing answer in answer. +- Do not call this tool while more tool work is still needed."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "answer": { + "type": "string", + "description": "Complete final answer to present to the user." + } + }, + "required": ["answer"] + }) +} diff --git a/crates/aish-tools/src/fs.rs b/crates/aish-tools/src/fs.rs deleted file mode 100644 index 49e63ab3..00000000 --- a/crates/aish-tools/src/fs.rs +++ /dev/null @@ -1,551 +0,0 @@ -use std::path::Path; - -use aish_i18n; -use aish_llm::{Tool, ToolResult}; - -/// Cached translated descriptions. -static READ_DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); -static WRITE_DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); -static EDIT_DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_read_description() -> &'static str { - READ_DESCRIPTION.get_or_init(|| aish_i18n::t("tools.fs.read_file.description")) -} - -fn get_write_description() -> &'static str { - WRITE_DESCRIPTION.get_or_init(|| aish_i18n::t("tools.fs.write_file.description")) -} - -fn get_edit_description() -> &'static str { - EDIT_DESCRIPTION.get_or_init(|| aish_i18n::t("tools.fs.edit_file.description")) -} - -/// Read file content tool. -pub struct ReadFileTool; - -impl Default for ReadFileTool { - fn default() -> Self { - Self::new() - } -} - -impl ReadFileTool { - pub fn new() -> Self { - Self - } -} - -impl Tool for ReadFileTool { - fn name(&self) -> &str { - "read_file" - } - - fn description(&self) -> &str { - get_read_description() - } - - fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "path": { "type": "string", "description": "Path to the file to read" }, - "offset": { "type": "integer", "description": "Line offset to start reading from (0-based)" }, - "limit": { "type": "integer", "description": "Maximum number of lines to read" } - }, - "required": ["path"] - }) - } - - fn execute(&self, args: serde_json::Value) -> ToolResult { - let path = match args.get("path").and_then(|p| p.as_str()) { - Some(p) => p, - None => return ToolResult::error(aish_i18n::t("tools.fs.read_file.missing_path")), - }; - - // Read raw bytes first for size check - let raw_bytes = match std::fs::read(path) { - Ok(b) => b, - Err(e) => { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - args_map.insert("error".to_string(), e.to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.read_file.read_failed", - &args_map, - )); - } - }; - - // Enforce 32KB size limit - const SIZE_LIMIT: usize = 32 * 1024; - if raw_bytes.len() > SIZE_LIMIT { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - args_map.insert("size".to_string(), raw_bytes.len().to_string()); - args_map.insert("limit".to_string(), SIZE_LIMIT.to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.read_file.file_too_large", - &args_map, - )); - } - - // Convert to UTF-8 - let content = match String::from_utf8(raw_bytes) { - Ok(s) => s, - Err(e) => { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - args_map.insert("error".to_string(), e.to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.read_file.decode_failed", - &args_map, - )); - } - }; - - let lines: Vec<&str> = content.lines().collect(); - - // Handle empty file - if lines.is_empty() { - return ToolResult::success(aish_i18n::t("tools.fs.read_file.empty_file")); - } - - let offset = args.get("offset").and_then(|o| o.as_u64()).unwrap_or(0) as usize; - let limit = args - .get("limit") - .and_then(|l| l.as_u64()) - .map(|l| l as usize); - - if offset >= lines.len() { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("offset".to_string(), offset.to_string()); - args_map.insert("length".to_string(), lines.len().to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.read_file.offset_exceeds_length", - &args_map, - )); - } - - // Format output with line numbers (1-based, offset-aware) - let selected: Vec = if let Some(limit) = limit { - lines - .iter() - .skip(offset) - .take(limit) - .enumerate() - .map(|(i, line)| format!("{:>6}\t{}", offset + i + 1, line)) - .collect() - } else { - lines - .iter() - .skip(offset) - .enumerate() - .map(|(i, line)| format!("{:>6}\t{}", offset + i + 1, line)) - .collect() - }; - - ToolResult::success(selected.join("\n")) - } -} - -/// Path-restricted wrapper around [`ReadFileTool`] for SSH sessions. -/// -/// Only allows reading files under the system temp directory's -/// `aish-offload/` subdirectory. Rejects any path outside that prefix -/// to prevent a remote LLM from reading arbitrary local files. -pub struct SshReadFileTool { - inner: ReadFileTool, - /// Canonicalized offload root (e.g. `/tmp/aish-offload`). - offload_root: std::path::PathBuf, -} - -impl SshReadFileTool { - pub fn new() -> Self { - let offload_root = std::env::temp_dir().join("aish-offload"); - // Best-effort canonicalize; if the dir doesn't exist yet, use - // the non-canonical form (first offload will create it). - let canonical_root = std::fs::canonicalize(&offload_root).unwrap_or(offload_root); - Self { - inner: ReadFileTool::new(), - offload_root: canonical_root, - } - } -} - -impl Tool for SshReadFileTool { - fn name(&self) -> &str { - self.inner.name() - } - - fn description(&self) -> &str { - self.inner.description() - } - - fn parameters(&self) -> serde_json::Value { - self.inner.parameters() - } - - fn execute(&self, args: serde_json::Value) -> ToolResult { - let path = match args.get("path").and_then(|p| p.as_str()) { - Some(p) => p, - None => return ToolResult::error(aish_i18n::t("tools.fs.read_file.missing_path")), - }; - // Canonicalize to resolve symlinks and '..' traversal. - let canonical = match std::fs::canonicalize(path) { - Ok(c) => c, - Err(e) => { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - args_map.insert("error".to_string(), e.to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.read_file.read_failed", - &args_map, - )); - } - }; - // Enforce exact offload root boundary. - if !canonical.starts_with(&self.offload_root) { - return ToolResult::error("Access denied: path is not inside offload directory"); - } - // Use the validated canonical path to avoid TOCTOU. - let mut safe_args = args; - if let Some(obj) = safe_args.as_object_mut() { - obj.insert( - "path".to_string(), - serde_json::Value::String(canonical.to_string_lossy().into_owned()), - ); - } - self.inner.execute(safe_args) - } -} - -/// Write file tool (creates or overwrites). -pub struct WriteFileTool; - -impl Default for WriteFileTool { - fn default() -> Self { - Self::new() - } -} - -impl WriteFileTool { - pub fn new() -> Self { - Self - } -} - -impl Tool for WriteFileTool { - fn name(&self) -> &str { - "write_file" - } - - fn description(&self) -> &str { - get_write_description() - } - - fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "path": { "type": "string", "description": "Path to the file to write" }, - "content": { "type": "string", "description": "Content to write" } - }, - "required": ["path", "content"] - }) - } - - fn execute(&self, args: serde_json::Value) -> ToolResult { - let path = match args.get("path").and_then(|p| p.as_str()) { - Some(p) => p, - None => return ToolResult::error(aish_i18n::t("tools.fs.write_file.missing_path")), - }; - let content = match args.get("content").and_then(|c| c.as_str()) { - Some(c) => c, - None => return ToolResult::error(aish_i18n::t("tools.fs.write_file.missing_content")), - }; - // Create parent dirs if needed - if let Some(parent) = Path::new(path).parent() { - if !parent.as_os_str().is_empty() { - if let Err(e) = std::fs::create_dir_all(parent) { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("error".to_string(), e.to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.write_file.create_dirs_failed", - &args_map, - )); - } - } - } - match std::fs::write(path, content) { - Ok(()) => { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("bytes".to_string(), content.len().to_string()); - args_map.insert("path".to_string(), path.to_string()); - ToolResult::success(aish_i18n::t_with_args( - "tools.fs.write_file.write_success", - &args_map, - )) - } - Err(e) => { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - args_map.insert("error".to_string(), e.to_string()); - ToolResult::error(aish_i18n::t_with_args( - "tools.fs.write_file.write_failed", - &args_map, - )) - } - } - } -} - -/// Edit file tool (string replacement). -pub struct EditFileTool; - -impl Default for EditFileTool { - fn default() -> Self { - Self::new() - } -} - -impl EditFileTool { - pub fn new() -> Self { - Self - } -} - -impl Tool for EditFileTool { - fn name(&self) -> &str { - "edit_file" - } - - fn description(&self) -> &str { - get_edit_description() - } - - fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "path": { "type": "string", "description": "Path to the file" }, - "old_string": { "type": "string", "description": "The text to replace" }, - "new_string": { "type": "string", "description": "The replacement text" }, - "replace_all": { "type": "boolean", "description": "Replace all occurrences (default: false)" } - }, - "required": ["path", "old_string", "new_string"] - }) - } - - fn execute(&self, args: serde_json::Value) -> ToolResult { - let path = match args.get("path").and_then(|p| p.as_str()) { - Some(p) => p, - None => return ToolResult::error(aish_i18n::t("tools.fs.edit_file.missing_path")), - }; - let old = match args.get("old_string").and_then(|o| o.as_str()) { - Some(o) => o, - None => { - return ToolResult::error(aish_i18n::t("tools.fs.edit_file.missing_old_string")) - } - }; - let new = match args.get("new_string").and_then(|n| n.as_str()) { - Some(n) => n, - None => { - return ToolResult::error(aish_i18n::t("tools.fs.edit_file.missing_new_string")) - } - }; - let replace_all = args - .get("replace_all") - .and_then(|r| r.as_bool()) - .unwrap_or(false); - - let content = match std::fs::read_to_string(path) { - Ok(c) => c, - Err(e) => { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - args_map.insert("error".to_string(), e.to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.edit_file.edit_read_failed", - &args_map, - )); - } - }; - - if !content.contains(old) { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.edit_file.old_string_not_found", - &args_map, - )); - } - - let new_content = if replace_all { - content.replace(old, new) - } else { - // Check uniqueness - let count = content.matches(old).count(); - if count > 1 { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("count".to_string(), count.to_string()); - args_map.insert("path".to_string(), path.to_string()); - return ToolResult::error(aish_i18n::t_with_args( - "tools.fs.edit_file.old_string_ambiguous", - &args_map, - )); - } - content.replacen(old, new, 1) - }; - - match std::fs::write(path, new_content) { - Ok(()) => { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - ToolResult::success(aish_i18n::t_with_args( - "tools.fs.edit_file.edit_success", - &args_map, - )) - } - Err(e) => { - let mut args_map = std::collections::HashMap::new(); - args_map.insert("path".to_string(), path.to_string()); - args_map.insert("error".to_string(), e.to_string()); - ToolResult::error(aish_i18n::t_with_args( - "tools.fs.edit_file.edit_write_failed", - &args_map, - )) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use aish_llm::Tool; - use std::fs; - - fn temp_dir() -> tempfile::TempDir { - tempfile::tempdir().expect("failed to create temp dir") - } - - #[test] - fn test_read_file_with_line_numbers() { - let dir = temp_dir(); - let file_path = dir.path().join("test.txt"); - fs::write(&file_path, "hello\nworld\nfoo").unwrap(); - - let tool = ReadFileTool::new(); - let result = tool.execute(serde_json::json!({ - "path": file_path.to_str().unwrap() - })); - - assert!(result.ok); - assert_eq!(result.output, " 1\thello\n 2\tworld\n 3\tfoo"); - } - - #[test] - fn test_read_file_with_offset() { - let dir = temp_dir(); - let file_path = dir.path().join("test.txt"); - fs::write(&file_path, "line1\nline2\nline3\nline4\nline5").unwrap(); - - let tool = ReadFileTool::new(); - let result = tool.execute(serde_json::json!({ - "path": file_path.to_str().unwrap(), - "offset": 2, - "limit": 2 - })); - - assert!(result.ok); - // Offset is 0-based, so offset=2 starts at line3 (3rd line) - // Line numbers are 1-based and offset-aware: 3, 4 - assert_eq!(result.output, " 3\tline3\n 4\tline4"); - } - - #[test] - fn test_read_file_size_limit() { - // Initialize i18n for testing - aish_i18n::set_locale("en-US"); - - let dir = temp_dir(); - let file_path = dir.path().join("big.txt"); - // Create a file larger than 32KB (33 * 1024 = 33792 bytes) - let big_content = "x".repeat(33 * 1024); - fs::write(&file_path, &big_content).unwrap(); - - let tool = ReadFileTool::new(); - let result = tool.execute(serde_json::json!({ - "path": file_path.to_str().unwrap() - })); - - assert!(!result.ok); - assert!( - result.output.contains("limit") || result.output.contains("bytes"), - "Expected size limit error, got: {}", - result.output - ); - } - - #[test] - fn test_edit_file_replace_all() { - let dir = temp_dir(); - let file_path = dir.path().join("test.txt"); - fs::write(&file_path, "foo bar foo baz foo").unwrap(); - - let tool = EditFileTool::new(); - let result = tool.execute(serde_json::json!({ - "path": file_path.to_str().unwrap(), - "old_string": "foo", - "new_string": "qux", - "replace_all": true - })); - - assert!(result.ok); - let content = fs::read_to_string(&file_path).unwrap(); - assert_eq!(content, "qux bar qux baz qux"); - } - - #[test] - fn test_edit_file_uniqueness_check() { - // Initialize i18n for testing - aish_i18n::set_locale("en-US"); - - let dir = temp_dir(); - let file_path = dir.path().join("test.txt"); - fs::write(&file_path, "foo bar foo baz").unwrap(); - - let tool = EditFileTool::new(); - let result = tool.execute(serde_json::json!({ - "path": file_path.to_str().unwrap(), - "old_string": "foo", - "new_string": "qux" - })); - - assert!(!result.ok); - assert!( - result.output.contains("times") || result.output.contains("ambiguous"), - "Expected uniqueness error, got: {}", - result.output - ); - // Verify file was NOT modified - let content = fs::read_to_string(&file_path).unwrap(); - assert_eq!(content, "foo bar foo baz"); - } - - #[test] - fn test_write_file_creates_parent_dirs() { - let dir = temp_dir(); - let file_path = dir.path().join("nested").join("deep").join("test.txt"); - - let tool = WriteFileTool::new(); - let result = tool.execute(serde_json::json!({ - "path": file_path.to_str().unwrap(), - "content": "hello world" - })); - - assert!(result.ok); - let content = fs::read_to_string(&file_path).unwrap(); - assert_eq!(content, "hello world"); - } -} diff --git a/crates/aish-tools/src/glob_tool.rs b/crates/aish-tools/src/glob_tool/glob_tool.rs similarity index 85% rename from crates/aish-tools/src/glob_tool.rs rename to crates/aish-tools/src/glob_tool/glob_tool.rs index f90c92f2..55974ea5 100644 --- a/crates/aish-tools/src/glob_tool.rs +++ b/crates/aish-tools/src/glob_tool/glob_tool.rs @@ -3,12 +3,7 @@ use std::path::PathBuf; use aish_i18n; use aish_llm::{Tool, ToolResult}; -/// Cached translated description. -static DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.glob.description")) -} +use super::prompt; /// Directories excluded by default (VCS and common large generated trees). const DEFAULT_EXCLUDE_DIRS: &[&str] = &[ @@ -48,24 +43,15 @@ impl Tool for GlobTool { } fn description(&self) -> &str { - get_description() + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "pattern": { - "type": "string", - "description": "Glob pattern such as **/*.py or src/**/*.md" - }, - "root": { - "type": "string", - "description": "Optional search root directory. Defaults to the current working directory." - } - }, - "required": ["pattern"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/glob_tool/prompt.rs b/crates/aish-tools/src/glob_tool/prompt.rs new file mode 100644 index 00000000..0a45cb4d --- /dev/null +++ b/crates/aish-tools/src/glob_tool/prompt.rs @@ -0,0 +1,25 @@ +pub(crate) const DESCRIPTION: &str = "Find files matching glob patterns."; + +pub(crate) const PROMPT: &str = r#"Use this tool to enumerate file paths by glob pattern. + +Usage: +- Prefer this tool when you need matching file names, not file contents. +- Use recursive patterns such as **/*.rs when searching across a tree. +- Use the root parameter to limit search scope when the user names a directory."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "Glob pattern such as **/*.py or src/**/*.md." + }, + "root": { + "type": "string", + "description": "Optional search root directory. Defaults to the current working directory." + } + }, + "required": ["pattern"] + }) +} diff --git a/crates/aish-tools/src/grep_tool.rs b/crates/aish-tools/src/grep_tool/grep_tool.rs similarity index 87% rename from crates/aish-tools/src/grep_tool.rs rename to crates/aish-tools/src/grep_tool/grep_tool.rs index 9e68b08f..694b0b3d 100644 --- a/crates/aish-tools/src/grep_tool.rs +++ b/crates/aish-tools/src/grep_tool/grep_tool.rs @@ -5,12 +5,7 @@ use std::path::PathBuf; use aish_i18n; use aish_llm::{Tool, ToolResult}; -/// Cached translated description. -static DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.grep.description")) -} +use super::prompt; /// Directories excluded by default (shared with GlobTool). const DEFAULT_EXCLUDE_DIRS: &[&str] = &[ @@ -51,28 +46,15 @@ impl Tool for GrepTool { } fn description(&self) -> &str { - get_description() + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "pattern": { - "type": "string", - "description": "Regex pattern to search for" - }, - "root": { - "type": "string", - "description": "Optional search root directory. Defaults to the current working directory." - }, - "include": { - "type": "string", - "description": "Optional glob filter for file names, e.g. *.py or *.rs" - } - }, - "required": ["pattern"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/grep_tool/prompt.rs b/crates/aish-tools/src/grep_tool/prompt.rs new file mode 100644 index 00000000..dd0c3ca1 --- /dev/null +++ b/crates/aish-tools/src/grep_tool/prompt.rs @@ -0,0 +1,30 @@ +pub(crate) const DESCRIPTION: &str = "Search file contents using a regex pattern."; + +pub(crate) const PROMPT: &str = r#"Use this tool to search text inside files. + +Usage: +- Use regex patterns for content search. +- Use root to limit the directory being searched. +- Use include to restrict matches to file names such as *.rs or *.py. +- Use glob when you only need matching file paths."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "Regex pattern to search for." + }, + "root": { + "type": "string", + "description": "Optional search root directory. Defaults to the current working directory." + }, + "include": { + "type": "string", + "description": "Optional glob filter for file names, e.g. *.py or *.rs." + } + }, + "required": ["pattern"] + }) +} diff --git a/crates/aish-tools/src/host_note.rs b/crates/aish-tools/src/host_note/host_note.rs similarity index 76% rename from crates/aish-tools/src/host_note.rs rename to crates/aish-tools/src/host_note/host_note.rs index 9ebe8624..ab7c2d02 100644 --- a/crates/aish-tools/src/host_note.rs +++ b/crates/aish-tools/src/host_note/host_note.rs @@ -5,11 +5,7 @@ use aish_llm::{Tool, ToolResult}; -static DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.host_note.description")) -} +use super::prompt; pub type HostNoteStoreFn = Box String + Send + Sync>; pub type HostNoteListFn = Box Vec + Send + Sync>; @@ -43,29 +39,15 @@ impl Tool for HostNoteTool { } fn description(&self) -> &str { - get_description() + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "action": { - "type": "string", - "enum": ["store", "list", "forget"], - "description": aish_i18n::t("tools.host_note.param.action") - }, - "content": { - "type": "string", - "description": aish_i18n::t("tools.host_note.param.content") - }, - "keyword": { - "type": "string", - "description": aish_i18n::t("tools.host_note.param.keyword") - } - }, - "required": ["action"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/host_note/prompt.rs b/crates/aish-tools/src/host_note/prompt.rs new file mode 100644 index 00000000..bb1a484b --- /dev/null +++ b/crates/aish-tools/src/host_note/prompt.rs @@ -0,0 +1,30 @@ +pub(crate) const DESCRIPTION: &str = "Store, list, or forget notes about the current remote host."; + +pub(crate) const PROMPT: &str = r#"Use this tool for notes scoped to the current remote host. + +Usage: +- Use store when the user provides durable facts about the host, services, paths, issues, or operational conventions. +- Use list before acting when host-specific context may matter. +- Use forget with a keyword to remove stale or incorrect notes."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "action": { + "type": "string", + "enum": ["store", "list", "forget"], + "description": "Host note operation to perform." + }, + "content": { + "type": "string", + "description": "Note content to save for the store action." + }, + "keyword": { + "type": "string", + "description": "Keyword used to match notes for the forget action." + } + }, + "required": ["action"] + }) +} diff --git a/crates/aish-tools/src/lib.rs b/crates/aish-tools/src/lib.rs index 4eacf716..bb2b6861 100644 --- a/crates/aish-tools/src/lib.rs +++ b/crates/aish-tools/src/lib.rs @@ -13,23 +13,144 @@ clippy::too_many_arguments )] -pub mod ask_user; -pub mod bash; -pub mod channel_ask_user; -pub mod channel_bash; -pub mod final_answer; -pub mod fs; -pub mod glob_tool; -pub mod grep_tool; -pub mod host_note; -pub mod memory_tool; -pub mod plan_tool; -pub mod python; pub mod registry; -pub mod secure_bash; -pub mod skill_tool; -pub mod system_diagnose; -pub mod web_fetch; + +pub mod ask_user { + mod ask_user; + mod prompt; + + pub use self::ask_user::*; +} + +pub mod bash { + mod bash; + mod prompt; + + pub use self::bash::*; +} + +pub mod channel_ask_user { + mod channel_ask_user; + mod prompt; + + pub use self::channel_ask_user::*; +} + +pub mod channel_bash { + mod channel_bash; + mod prompt; + + pub use self::channel_bash::*; +} + +pub mod final_answer { + mod final_answer; + mod prompt; + + pub use self::final_answer::*; +} + +pub mod edit_file { + mod edit_file; + mod prompt; + + pub use self::edit_file::*; +} + +pub mod fs { + pub use crate::edit_file::EditFileTool; + pub use crate::read_file::{ReadFileTool, SshReadFileTool}; + pub use crate::write_file::WriteFileTool; +} + +pub mod glob_tool { + mod glob_tool; + mod prompt; + + pub use self::glob_tool::*; +} + +pub mod grep_tool { + mod grep_tool; + mod prompt; + + pub use self::grep_tool::*; +} + +pub mod host_note { + mod host_note; + mod prompt; + + pub use self::host_note::*; +} + +pub mod memory_tool { + mod memory_tool; + mod prompt; + + pub use self::memory_tool::*; +} + +pub mod plan_tool { + mod enter_plan_mode; + mod exit_plan_mode; + mod list_plan_templates; + mod prompt; + + pub use self::enter_plan_mode::EnterPlanModeTool; + pub use self::exit_plan_mode::ExitPlanModeTool; + pub use self::list_plan_templates::ListTemplatesTool; +} + +pub mod python { + mod prompt; + mod python; + + pub use self::python::*; +} + +pub mod read_file { + mod prompt; + mod read_file; + + pub use self::read_file::*; +} + +pub mod secure_bash { + mod secure_bash; + + pub use self::secure_bash::*; +} + +pub mod skill_tool { + mod prompt; + mod skill_tool; + + pub use self::skill_tool::*; +} + +pub mod system_diagnose { + mod prompt; + mod system_diagnose; + + pub use self::system_diagnose::*; +} + +pub mod web_fetch { + mod preapproved; + mod prompt; + mod utils; + mod web_fetch; + + pub use self::web_fetch::*; +} + +pub mod write_file { + mod prompt; + mod write_file; + + pub use self::write_file::*; +} pub use ask_user::AskUserTool; pub use channel_ask_user::ChannelAskUserTool; diff --git a/crates/aish-tools/src/memory_tool.rs b/crates/aish-tools/src/memory_tool/memory_tool.rs similarity index 79% rename from crates/aish-tools/src/memory_tool.rs rename to crates/aish-tools/src/memory_tool/memory_tool.rs index 08874cd4..e1d052df 100644 --- a/crates/aish-tools/src/memory_tool.rs +++ b/crates/aish-tools/src/memory_tool/memory_tool.rs @@ -1,12 +1,7 @@ use aish_i18n; use aish_llm::{Tool, ToolResult}; -/// Cached translated description. -static DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.memory.description")) -} +use super::prompt; /// Callback type for memory operations. pub type MemorySearchFn = Box Vec + Send + Sync>; @@ -62,38 +57,15 @@ impl Tool for MemoryTool { } fn description(&self) -> &str { - get_description() + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "action": { - "type": "string", - "enum": ["search", "store", "forget", "list"], - "description": "Memory operation to perform" - }, - "query": { - "type": "string", - "description": "Search query (for 'search' action)" - }, - "content": { - "type": "string", - "description": "Content to store (for 'store' action)" - }, - "category": { - "type": "string", - "enum": ["preference", "environment", "solution", "pattern", "other"], - "description": "Category for stored memory (default: other)" - }, - "memory_id": { - "type": "integer", - "description": "Memory ID to forget (for 'forget' action)" - } - }, - "required": ["action"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/memory_tool/prompt.rs b/crates/aish-tools/src/memory_tool/prompt.rs new file mode 100644 index 00000000..1da2053e --- /dev/null +++ b/crates/aish-tools/src/memory_tool/prompt.rs @@ -0,0 +1,40 @@ +pub(crate) const DESCRIPTION: &str = "Search, store, list, or forget long-term memories."; + +pub(crate) const PROMPT: &str = r#"Use this tool for long-term memory operations. + +Usage: +- Use search to recall relevant saved knowledge before acting on user preferences or prior context. +- Use store only for durable facts that are likely useful later. +- Use list to inspect recent memories. +- Use forget to remove outdated or incorrect memory entries by id."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "action": { + "type": "string", + "enum": ["search", "store", "forget", "list"], + "description": "Memory operation to perform." + }, + "query": { + "type": "string", + "description": "Search query for the search action." + }, + "content": { + "type": "string", + "description": "Content to store for the store action." + }, + "category": { + "type": "string", + "enum": ["preference", "environment", "solution", "pattern", "other"], + "description": "Category for stored memory. Defaults to other." + }, + "memory_id": { + "type": "integer", + "description": "Memory id to forget for the forget action." + } + }, + "required": ["action"] + }) +} diff --git a/crates/aish-tools/src/plan_tool.rs b/crates/aish-tools/src/plan_tool.rs deleted file mode 100644 index 995a2fb2..00000000 --- a/crates/aish-tools/src/plan_tool.rs +++ /dev/null @@ -1,488 +0,0 @@ -// Plan mode tools for entering and exiting structured planning phase. - -use aish_core::plan::generate_plan_id; -use aish_llm::{Tool, ToolResult}; - -/// Tools visible during planning phase (mirrors aish_core::plan::PLANNING_VISIBLE_TOOLS). -const VISIBLE_TOOLS_DURING_PLANNING: &[&str] = &[ - "read_file", - "glob", - "grep", - "ask_user", - "memory", - "write_file", - "edit_file", - "exit_plan_mode", -]; - -/// Tool for entering plan mode. -/// -/// When the AI calls this tool, it transitions to a planning phase where -/// only read-only tools plus write_file/edit_file (for the plan artifact) are available. -/// The tool returns metadata that the session layer uses to initialize plan state. -pub struct EnterPlanModeTool; - -impl EnterPlanModeTool { - pub fn new() -> Self { - Self - } -} - -impl Tool for EnterPlanModeTool { - fn name(&self) -> &str { - "enter_plan_mode" - } - - fn description(&self) -> &str { - "Enter plan mode to design an implementation approach before writing code. \ - During planning, only read-only tools and write_file/edit_file (for the plan) are available. \ - When ready, use exit_plan_mode to present the plan for approval." - } - - fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "topic": { - "type": "string", - "description": "The topic or task to plan" - }, - "summary": { - "type": "string", - "description": "Brief summary of the planning goal (optional)" - } - }, - "required": ["topic"] - }) - } - - fn execute(&self, args: serde_json::Value) -> ToolResult { - // Extract arguments - let topic = args - .get("topic") - .and_then(|v| v.as_str()) - .unwrap_or("unknown"); - - let summary = args - .get("summary") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - // Generate a plan ID for this planning session - let plan_id = generate_plan_id(); - - // Build suggested artifact path (session layer will finalize the actual path) - let artifact_suggestion = format!(".aish/plans/plan-{}.md", plan_id); - - // Visible tools list for the AI's reference - let visible_tools: Vec<&str> = VISIBLE_TOOLS_DURING_PLANNING.to_vec(); - - // Return metadata for the session layer to initialize plan state. - let meta = serde_json::json!({ - "action": "enter_plan_mode", - "topic": topic, - "summary": summary, - "plan_id": plan_id, - "phase": "Planning", - "visible_tools": visible_tools, - "artifact_suggestion": artifact_suggestion - }); - - ToolResult { - ok: true, - output: format!( - "Entering plan mode for: {}\n\ - Plan ID: {}\n\n\ - During planning, you have access to:\n\ - - Read-only tools: read_file, glob, grep, ask_user, memory\n\ - - Write tools (for plan only): write_file, edit_file\n\ - - exit_plan_mode: when ready to present your plan\n\n\ - Use write_file to create your plan artifact.\n\ - Suggested path: {}", - topic, plan_id, artifact_suggestion - ), - meta: Some(meta), - } - } -} - -/// Tool for exiting plan mode. -/// -/// When the AI calls this tool, it exits the planning phase and presents -/// the plan for user approval. The tool reads the plan artifact content -/// and returns it along with approval instructions. -/// -/// The actual approval interaction happens at the shell layer (app.rs), -/// not in the tool itself. -pub struct ExitPlanModeTool; - -impl ExitPlanModeTool { - pub fn new() -> Self { - Self - } -} - -impl Tool for ExitPlanModeTool { - fn name(&self) -> &str { - "exit_plan_mode" - } - - fn description(&self) -> &str { - "Exit plan mode and present your plan for approval. \ - The plan will be reviewed by the user before proceeding with implementation." - } - - fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "summary": { - "type": "string", - "description": "Brief summary of the plan (optional)" - }, - "feedback": { - "type": "string", - "description": "Feedback from user when changes are requested (optional, injected by session layer)" - }, - "plan_content": { - "type": "string", - "description": "Full plan content for review (optional, injected by session layer)" - } - } - }) - } - - fn execute(&self, args: serde_json::Value) -> ToolResult { - // Extract summary - let summary = args - .get("summary") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - // Extract feedback (injected by session layer when changes are requested) - let feedback = args - .get("feedback") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - // Extract plan content (injected by session layer) - let plan_content = args - .get("plan_content") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - // This tool doesn't directly read the artifact or manage state. - // It signals the session layer to handle the approval workflow. - // The session layer will: - // 1. Read the plan artifact - // 2. Present it to the user for approval via PlanApprovalFlow - // 3. Create an approved snapshot or relay feedback back to the AI - - // Build metadata with approval state transitions - let mut meta = serde_json::json!({ - "action": "exit_plan_mode", - "decision_required": true, - "summary": summary - }); - - // Include feedback if present (when re-exiting after changes requested) - if let Some(ref fb) = feedback { - meta["feedback"] = serde_json::json!(fb); - meta["approval_transition"] = serde_json::json!("changes_requested_to_review"); - } - - // Include plan content hint if present - if let Some(ref content) = plan_content { - meta["plan_content_length"] = serde_json::json!(content.len()); - } - - ToolResult { - ok: true, - output: "Plan mode exited. The plan is now ready for review and approval.".to_string(), - meta: Some(meta), - } - } -} - -/// Tool for listing available plan templates. -/// -/// Returns a list of structured templates the AI can use when creating -/// plan artifacts during plan mode. -pub struct ListTemplatesTool; - -impl ListTemplatesTool { - pub fn new() -> Self { - Self - } -} - -impl Tool for ListTemplatesTool { - fn name(&self) -> &str { - "list_plan_templates" - } - - fn description(&self) -> &str { - "List available plan templates for structuring your implementation plan." - } - - fn parameters(&self) -> serde_json::Value { - serde_json::json!({"type": "object", "properties": {}}) - } - - fn execute(&self, _args: serde_json::Value) -> ToolResult { - let templates = aish_core::plan::get_available_templates(); - let template_list: Vec = templates - .iter() - .map(|t| { - serde_json::json!({ - "name": t.name, - "description": t.description, - "content": t.content, - }) - }) - .collect(); - - let output = templates - .iter() - .map(|t| format!("- **{}**: {}", t.name, t.description)) - .collect::>() - .join("\n"); - - ToolResult { - ok: true, - output: format!("Available plan templates:\n{}", output), - meta: Some(serde_json::json!({ - "templates": template_list - })), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_enter_plan_mode_basic() { - let tool = EnterPlanModeTool::new(); - assert_eq!(tool.name(), "enter_plan_mode"); - - let result = tool.execute(serde_json::json!({ - "topic": "implement feature X" - })); - - assert!(result.ok); - assert!(result.output.contains("Entering plan mode")); - assert!(result.output.contains("implement feature X")); - - // Check metadata - assert!(result.meta.is_some()); - let meta = result.meta.unwrap(); - assert_eq!(meta["action"], "enter_plan_mode"); - assert_eq!(meta["topic"], "implement feature X"); - assert_eq!(meta["phase"], "Planning"); - - // New fields - assert!(meta["plan_id"].is_string()); - assert_eq!(meta["plan_id"].as_str().unwrap().len(), 12); - assert!(meta["visible_tools"].is_array()); - assert!(meta["artifact_suggestion"].is_string()); - } - - #[test] - fn test_enter_plan_mode_with_summary() { - let tool = EnterPlanModeTool::new(); - let result = tool.execute(serde_json::json!({ - "topic": "refactor code", - "summary": "Clean up module structure" - })); - - assert!(result.ok); - let meta = result.meta.unwrap(); - assert_eq!(meta["summary"], "Clean up module structure"); - } - - #[test] - fn test_enter_plan_mode_missing_topic() { - let tool = EnterPlanModeTool::new(); - let result = tool.execute(serde_json::json!({})); - - assert!(result.ok); // Should still succeed with default topic - assert!(result.output.contains("Entering plan mode")); - } - - #[test] - fn test_enter_plan_mode_unique_plan_ids() { - let tool = EnterPlanModeTool::new(); - let r1 = tool.execute(serde_json::json!({"topic": "a"})); - let r2 = tool.execute(serde_json::json!({"topic": "b"})); - - let id1 = r1.meta.as_ref().unwrap()["plan_id"].as_str().unwrap(); - let id2 = r2.meta.as_ref().unwrap()["plan_id"].as_str().unwrap(); - assert_ne!(id1, id2); - } - - #[test] - fn test_enter_plan_mode_visible_tools() { - let tool = EnterPlanModeTool::new(); - let result = tool.execute(serde_json::json!({"topic": "test"})); - - let meta = result.meta.unwrap(); - let visible: Vec<&str> = meta["visible_tools"] - .as_array() - .unwrap() - .iter() - .map(|v| v.as_str().unwrap()) - .collect(); - - assert!(visible.contains(&"read_file")); - assert!(visible.contains(&"write_file")); - assert!(visible.contains(&"edit_file")); - assert!(visible.contains(&"exit_plan_mode")); - assert!(!visible.contains(&"bash_exec")); - } - - #[test] - fn test_exit_plan_mode_basic() { - let tool = ExitPlanModeTool::new(); - assert_eq!(tool.name(), "exit_plan_mode"); - - let result = tool.execute(serde_json::json!({})); - - assert!(result.ok); - assert!(result.output.contains("Plan mode exited")); - assert!(result.output.contains("ready for review")); - - // Check metadata - assert!(result.meta.is_some()); - let meta = result.meta.unwrap(); - assert_eq!(meta["action"], "exit_plan_mode"); - assert_eq!(meta["decision_required"], true); - } - - #[test] - fn test_exit_plan_mode_with_summary() { - let tool = ExitPlanModeTool::new(); - let result = tool.execute(serde_json::json!({ - "summary": "Complete implementation plan" - })); - - assert!(result.ok); - let meta = result.meta.unwrap(); - assert_eq!(meta["summary"], "Complete implementation plan"); - } - - #[test] - fn test_tool_descriptions() { - let enter = EnterPlanModeTool::new(); - let exit = ExitPlanModeTool::new(); - let templates = ListTemplatesTool::new(); - - assert!(enter.description().contains("plan mode")); - assert!(enter.description().contains("read-only")); - - assert!(exit.description().contains("approval")); - assert!(exit.description().contains("review")); - - assert!(templates.description().contains("templates")); - } - - #[test] - fn test_enter_plan_mode_parameters() { - let tool = EnterPlanModeTool::new(); - let params = tool.parameters(); - - assert_eq!(params["type"], "object"); - assert!(params["properties"]["topic"]["description"] - .as_str() - .is_some()); - assert!(params["properties"]["summary"]["description"] - .as_str() - .is_some()); - - let required = params["required"].as_array().unwrap(); - assert_eq!(required.len(), 1); - assert_eq!(required[0], "topic"); - } - - #[test] - fn test_exit_plan_mode_parameters() { - let tool = ExitPlanModeTool::new(); - let params = tool.parameters(); - - assert_eq!(params["type"], "object"); - assert!(params["properties"]["summary"]["description"] - .as_str() - .is_some()); - assert!(params["properties"]["feedback"]["description"] - .as_str() - .is_some()); - assert!(params["properties"]["plan_content"]["description"] - .as_str() - .is_some()); - - // No required parameters - let required = params["required"].as_array(); - assert!(required.is_none() || required.unwrap().is_empty()); - } - - #[test] - fn test_exit_plan_mode_with_feedback() { - let tool = ExitPlanModeTool::new(); - let result = tool.execute(serde_json::json!({ - "summary": "Revised plan", - "feedback": "Please add more testing steps" - })); - - assert!(result.ok); - let meta = result.meta.unwrap(); - assert_eq!(meta["feedback"], "Please add more testing steps"); - assert_eq!(meta["approval_transition"], "changes_requested_to_review"); - } - - #[test] - fn test_exit_plan_mode_with_plan_content() { - let tool = ExitPlanModeTool::new(); - let result = tool.execute(serde_json::json!({ - "summary": "Full plan", - "plan_content": "# Plan\n## Steps\n1. Do stuff\n2. Test" - })); - - assert!(result.ok); - let meta = result.meta.unwrap(); - assert!(meta["plan_content_length"].is_number()); - } - - #[test] - fn test_list_templates_tool() { - let tool = ListTemplatesTool::new(); - assert_eq!(tool.name(), "list_plan_templates"); - - let result = tool.execute(serde_json::json!({})); - assert!(result.ok); - assert!(result.output.contains("Available plan templates")); - assert!(result.output.contains("default")); - assert!(result.output.contains("bugfix")); - assert!(result.output.contains("feature")); - - let meta = result.meta.unwrap(); - let templates = meta["templates"].as_array().unwrap(); - assert_eq!(templates.len(), 3); - - // Check that each template has required fields - for t in templates { - assert!(t["name"].is_string()); - assert!(t["description"].is_string()); - assert!(t["content"].is_string()); - } - } - - #[test] - fn test_list_templates_parameters() { - let tool = ListTemplatesTool::new(); - let params = tool.parameters(); - assert_eq!(params["type"], "object"); - // No properties needed - assert!(params["properties"].as_object().unwrap().is_empty()); - } -} diff --git a/crates/aish-tools/src/plan_tool/enter_plan_mode.rs b/crates/aish-tools/src/plan_tool/enter_plan_mode.rs new file mode 100644 index 00000000..79a86bcd --- /dev/null +++ b/crates/aish-tools/src/plan_tool/enter_plan_mode.rs @@ -0,0 +1,191 @@ +use aish_core::plan::generate_plan_id; +use aish_llm::{Tool, ToolResult}; + +use super::prompt; + +/// Tools visible during planning phase (mirrors aish_core::plan::PLANNING_VISIBLE_TOOLS). +const VISIBLE_TOOLS_DURING_PLANNING: &[&str] = &[ + "read_file", + "glob", + "grep", + "ask_user", + "memory", + "write_file", + "edit_file", + "exit_plan_mode", +]; + +/// Tool for entering plan mode. +pub struct EnterPlanModeTool; + +impl EnterPlanModeTool { + pub fn new() -> Self { + Self + } +} + +impl Tool for EnterPlanModeTool { + fn name(&self) -> &str { + "enter_plan_mode" + } + + fn description(&self) -> &str { + prompt::ENTER_DESCRIPTION + } + + fn parameters(&self) -> serde_json::Value { + prompt::enter_parameters() + } + + fn prompt(&self) -> &str { + prompt::ENTER_PROMPT + } + + fn execute(&self, args: serde_json::Value) -> ToolResult { + let topic = args + .get("topic") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + + let summary = args + .get("summary") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let plan_id = generate_plan_id(); + let artifact_suggestion = format!(".aish/plans/plan-{}.md", plan_id); + let visible_tools: Vec<&str> = VISIBLE_TOOLS_DURING_PLANNING.to_vec(); + + let meta = serde_json::json!({ + "action": "enter_plan_mode", + "topic": topic, + "summary": summary, + "plan_id": plan_id, + "phase": "Planning", + "visible_tools": visible_tools, + "artifact_suggestion": artifact_suggestion + }); + + ToolResult { + ok: true, + output: format!( + "Entering plan mode for: {}\n\ + Plan ID: {}\n\n\ + During planning, you have access to:\n\ + - Read-only tools: read_file, glob, grep, ask_user, memory\n\ + - Write tools (for plan only): write_file, edit_file\n\ + - exit_plan_mode: when ready to present your plan\n\n\ + Use write_file to create your plan artifact.\n\ + Suggested path: {}", + topic, plan_id, artifact_suggestion + ), + meta: Some(meta), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_enter_plan_mode_basic() { + let tool = EnterPlanModeTool::new(); + assert_eq!(tool.name(), "enter_plan_mode"); + + let result = tool.execute(serde_json::json!({ + "topic": "implement feature X" + })); + + assert!(result.ok); + assert!(result.output.contains("Entering plan mode")); + assert!(result.output.contains("implement feature X")); + + let meta = result.meta.unwrap(); + assert_eq!(meta["action"], "enter_plan_mode"); + assert_eq!(meta["topic"], "implement feature X"); + assert_eq!(meta["phase"], "Planning"); + assert!(meta["plan_id"].is_string()); + assert_eq!(meta["plan_id"].as_str().unwrap().len(), 12); + assert!(meta["visible_tools"].is_array()); + assert!(meta["artifact_suggestion"].is_string()); + } + + #[test] + fn test_enter_plan_mode_with_summary() { + let tool = EnterPlanModeTool::new(); + let result = tool.execute(serde_json::json!({ + "topic": "refactor code", + "summary": "Clean up module structure" + })); + + assert!(result.ok); + let meta = result.meta.unwrap(); + assert_eq!(meta["summary"], "Clean up module structure"); + } + + #[test] + fn test_enter_plan_mode_missing_topic() { + let tool = EnterPlanModeTool::new(); + let result = tool.execute(serde_json::json!({})); + + assert!(result.ok); + assert!(result.output.contains("Entering plan mode")); + } + + #[test] + fn test_enter_plan_mode_unique_plan_ids() { + let tool = EnterPlanModeTool::new(); + let r1 = tool.execute(serde_json::json!({"topic": "a"})); + let r2 = tool.execute(serde_json::json!({"topic": "b"})); + + let id1 = r1.meta.as_ref().unwrap()["plan_id"].as_str().unwrap(); + let id2 = r2.meta.as_ref().unwrap()["plan_id"].as_str().unwrap(); + assert_ne!(id1, id2); + } + + #[test] + fn test_enter_plan_mode_visible_tools() { + let tool = EnterPlanModeTool::new(); + let result = tool.execute(serde_json::json!({"topic": "test"})); + + let meta = result.meta.unwrap(); + let visible: Vec<&str> = meta["visible_tools"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap()) + .collect(); + + assert!(visible.contains(&"read_file")); + assert!(visible.contains(&"write_file")); + assert!(visible.contains(&"edit_file")); + assert!(visible.contains(&"exit_plan_mode")); + assert!(!visible.contains(&"bash_exec")); + } + + #[test] + fn test_enter_plan_mode_parameters() { + let tool = EnterPlanModeTool::new(); + let params = tool.parameters(); + + assert_eq!(params["type"], "object"); + assert!(params["properties"]["topic"]["description"] + .as_str() + .is_some()); + assert!(params["properties"]["summary"]["description"] + .as_str() + .is_some()); + + let required = params["required"].as_array().unwrap(); + assert_eq!(required.len(), 1); + assert_eq!(required[0], "topic"); + } + + #[test] + fn test_enter_plan_mode_description() { + let tool = EnterPlanModeTool::new(); + assert!(tool.description().contains("plan mode")); + assert!(tool.description().contains("read-only")); + } +} diff --git a/crates/aish-tools/src/plan_tool/exit_plan_mode.rs b/crates/aish-tools/src/plan_tool/exit_plan_mode.rs new file mode 100644 index 00000000..8713ab83 --- /dev/null +++ b/crates/aish-tools/src/plan_tool/exit_plan_mode.rs @@ -0,0 +1,155 @@ +use aish_llm::{Tool, ToolResult}; + +use super::prompt; + +/// Tool for exiting plan mode. +pub struct ExitPlanModeTool; + +impl ExitPlanModeTool { + pub fn new() -> Self { + Self + } +} + +impl Tool for ExitPlanModeTool { + fn name(&self) -> &str { + "exit_plan_mode" + } + + fn description(&self) -> &str { + prompt::EXIT_DESCRIPTION + } + + fn parameters(&self) -> serde_json::Value { + prompt::exit_parameters() + } + + fn prompt(&self) -> &str { + prompt::EXIT_PROMPT + } + + fn execute(&self, args: serde_json::Value) -> ToolResult { + let summary = args + .get("summary") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let feedback = args + .get("feedback") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let plan_content = args + .get("plan_content") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let mut meta = serde_json::json!({ + "action": "exit_plan_mode", + "decision_required": true, + "summary": summary + }); + + if let Some(ref fb) = feedback { + meta["feedback"] = serde_json::json!(fb); + meta["approval_transition"] = serde_json::json!("changes_requested_to_review"); + } + + if let Some(ref content) = plan_content { + meta["plan_content_length"] = serde_json::json!(content.len()); + } + + ToolResult { + ok: true, + output: "Plan mode exited. The plan is now ready for review and approval.".to_string(), + meta: Some(meta), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_exit_plan_mode_basic() { + let tool = ExitPlanModeTool::new(); + assert_eq!(tool.name(), "exit_plan_mode"); + + let result = tool.execute(serde_json::json!({})); + + assert!(result.ok); + assert!(result.output.contains("Plan mode exited")); + assert!(result.output.contains("ready for review")); + + let meta = result.meta.unwrap(); + assert_eq!(meta["action"], "exit_plan_mode"); + assert_eq!(meta["decision_required"], true); + } + + #[test] + fn test_exit_plan_mode_with_summary() { + let tool = ExitPlanModeTool::new(); + let result = tool.execute(serde_json::json!({ + "summary": "Complete implementation plan" + })); + + assert!(result.ok); + let meta = result.meta.unwrap(); + assert_eq!(meta["summary"], "Complete implementation plan"); + } + + #[test] + fn test_exit_plan_mode_parameters() { + let tool = ExitPlanModeTool::new(); + let params = tool.parameters(); + + assert_eq!(params["type"], "object"); + assert!(params["properties"]["summary"]["description"] + .as_str() + .is_some()); + assert!(params["properties"]["feedback"]["description"] + .as_str() + .is_some()); + assert!(params["properties"]["plan_content"]["description"] + .as_str() + .is_some()); + + let required = params["required"].as_array(); + assert!(required.is_none() || required.unwrap().is_empty()); + } + + #[test] + fn test_exit_plan_mode_with_feedback() { + let tool = ExitPlanModeTool::new(); + let result = tool.execute(serde_json::json!({ + "summary": "Revised plan", + "feedback": "Please add more testing steps" + })); + + assert!(result.ok); + let meta = result.meta.unwrap(); + assert_eq!(meta["feedback"], "Please add more testing steps"); + assert_eq!(meta["approval_transition"], "changes_requested_to_review"); + } + + #[test] + fn test_exit_plan_mode_with_plan_content() { + let tool = ExitPlanModeTool::new(); + let result = tool.execute(serde_json::json!({ + "summary": "Full plan", + "plan_content": "# Plan\n## Steps\n1. Do stuff\n2. Test" + })); + + assert!(result.ok); + let meta = result.meta.unwrap(); + assert!(meta["plan_content_length"].is_number()); + } + + #[test] + fn test_exit_plan_mode_description() { + let tool = ExitPlanModeTool::new(); + assert!(tool.description().contains("approval")); + assert!(tool.description().contains("review")); + } +} diff --git a/crates/aish-tools/src/plan_tool/list_plan_templates.rs b/crates/aish-tools/src/plan_tool/list_plan_templates.rs new file mode 100644 index 00000000..43c68d2e --- /dev/null +++ b/crates/aish-tools/src/plan_tool/list_plan_templates.rs @@ -0,0 +1,100 @@ +use aish_llm::{Tool, ToolResult}; + +use super::prompt; + +/// Tool for listing available plan templates. +pub struct ListTemplatesTool; + +impl ListTemplatesTool { + pub fn new() -> Self { + Self + } +} + +impl Tool for ListTemplatesTool { + fn name(&self) -> &str { + "list_plan_templates" + } + + fn description(&self) -> &str { + prompt::TEMPLATES_DESCRIPTION + } + + fn parameters(&self) -> serde_json::Value { + prompt::templates_parameters() + } + + fn prompt(&self) -> &str { + prompt::TEMPLATES_PROMPT + } + + fn execute(&self, _args: serde_json::Value) -> ToolResult { + let templates = aish_core::plan::get_available_templates(); + let template_list: Vec = templates + .iter() + .map(|t| { + serde_json::json!({ + "name": t.name, + "description": t.description, + "content": t.content, + }) + }) + .collect(); + + let output = templates + .iter() + .map(|t| format!("- **{}**: {}", t.name, t.description)) + .collect::>() + .join("\n"); + + ToolResult { + ok: true, + output: format!("Available plan templates:\n{}", output), + meta: Some(serde_json::json!({ + "templates": template_list + })), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_list_templates_tool() { + let tool = ListTemplatesTool::new(); + assert_eq!(tool.name(), "list_plan_templates"); + + let result = tool.execute(serde_json::json!({})); + assert!(result.ok); + assert!(result.output.contains("Available plan templates")); + assert!(result.output.contains("default")); + assert!(result.output.contains("bugfix")); + assert!(result.output.contains("feature")); + + let meta = result.meta.unwrap(); + let templates = meta["templates"].as_array().unwrap(); + assert_eq!(templates.len(), 3); + + for t in templates { + assert!(t["name"].is_string()); + assert!(t["description"].is_string()); + assert!(t["content"].is_string()); + } + } + + #[test] + fn test_list_templates_parameters() { + let tool = ListTemplatesTool::new(); + let params = tool.parameters(); + assert_eq!(params["type"], "object"); + assert!(params["properties"].as_object().unwrap().is_empty()); + } + + #[test] + fn test_list_templates_description() { + let tool = ListTemplatesTool::new(); + assert!(tool.description().contains("templates")); + } +} diff --git a/crates/aish-tools/src/plan_tool/prompt.rs b/crates/aish-tools/src/plan_tool/prompt.rs new file mode 100644 index 00000000..f2e81544 --- /dev/null +++ b/crates/aish-tools/src/plan_tool/prompt.rs @@ -0,0 +1,63 @@ +pub(crate) const ENTER_DESCRIPTION: &str = + "Enter plan mode to design an implementation plan with read-only planning tools."; +pub(crate) const EXIT_DESCRIPTION: &str = + "Exit plan mode and present the plan for approval and review."; +pub(crate) const TEMPLATES_DESCRIPTION: &str = "List available plan templates."; + +pub(crate) const ENTER_PROMPT: &str = r#"Use this tool when a task needs structured planning before implementation. + +Usage: +- Enter plan mode before making changes for multi-step or risky work. +- During planning, use read-only tools plus write_file/edit_file for the plan artifact. +- Exit plan mode when the plan is ready for user approval."#; + +pub(crate) const EXIT_PROMPT: &str = r#"Use this tool when the plan is ready for review. + +Usage: +- Ensure the plan artifact is complete before exiting plan mode. +- Include a concise summary when helpful. +- If feedback was provided, address it before re-submitting."#; + +pub(crate) const TEMPLATES_PROMPT: &str = + r#"Use this tool to inspect available plan templates before writing a plan artifact."#; + +pub(crate) fn enter_parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Topic or task to plan." + }, + "summary": { + "type": "string", + "description": "Optional brief summary of the planning goal." + } + }, + "required": ["topic"] + }) +} + +pub(crate) fn exit_parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "summary": { + "type": "string", + "description": "Optional brief summary of the plan." + }, + "feedback": { + "type": "string", + "description": "Optional feedback from the user when changes are requested. Injected by the session layer." + }, + "plan_content": { + "type": "string", + "description": "Optional full plan content for review. Injected by the session layer." + } + } + }) +} + +pub(crate) fn templates_parameters() -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) +} diff --git a/crates/aish-tools/src/python/prompt.rs b/crates/aish-tools/src/python/prompt.rs new file mode 100644 index 00000000..e32c0197 --- /dev/null +++ b/crates/aish-tools/src/python/prompt.rs @@ -0,0 +1,22 @@ +pub(crate) const DESCRIPTION: &str = "Execute Python code and return the result."; + +pub(crate) const PROMPT: &str = r#"Use this tool for small Python snippets that are better expressed as code than shell pipelines. + +Usage: +- Print values that should be returned to the conversation. +- Keep snippets focused and self-contained. +- Prefer this tool for structured data processing, calculations, and short scripts. +- Do not use this tool for long-running or interactive programs."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "Python code to execute." + } + }, + "required": ["code"] + }) +} diff --git a/crates/aish-tools/src/python.rs b/crates/aish-tools/src/python/python.rs similarity index 88% rename from crates/aish-tools/src/python.rs rename to crates/aish-tools/src/python/python.rs index 053bb8ab..f0d532db 100644 --- a/crates/aish-tools/src/python.rs +++ b/crates/aish-tools/src/python/python.rs @@ -3,16 +3,11 @@ use std::process::Command; use aish_i18n; use aish_llm::{Tool, ToolResult}; +use super::prompt; + /// Maximum output length (matches main branch's 1000 chars). const MAX_OUTPUT_CHARS: usize = 1000; -/// Cached translated description. -static DESCRIPTION: std::sync::OnceLock = std::sync::OnceLock::new(); - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.python.description")) -} - /// Tool for executing Python code. pub struct PythonTool; @@ -34,20 +29,15 @@ impl Tool for PythonTool { } fn description(&self) -> &str { - get_description() + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "code": { - "type": "string", - "description": "The Python code to execute." - } - }, - "required": ["code"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/read_file/prompt.rs b/crates/aish-tools/src/read_file/prompt.rs new file mode 100644 index 00000000..1b516abf --- /dev/null +++ b/crates/aish-tools/src/read_file/prompt.rs @@ -0,0 +1,29 @@ +pub(crate) const DESCRIPTION: &str = "Read text content from a file."; + +pub(crate) const PROMPT: &str = r#"Use this tool to read text files. + +Usage: +- Provide path to the file to read. +- Use offset and limit when you only need part of a larger file. +- Results include 1-based line numbers."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file to read." + }, + "offset": { + "type": "integer", + "description": "Line offset to start reading from, 0-based." + }, + "limit": { + "type": "integer", + "description": "Maximum number of lines to read." + } + }, + "required": ["path"] + }) +} diff --git a/crates/aish-tools/src/read_file/read_file.rs b/crates/aish-tools/src/read_file/read_file.rs new file mode 100644 index 00000000..115dddcc --- /dev/null +++ b/crates/aish-tools/src/read_file/read_file.rs @@ -0,0 +1,252 @@ +use aish_i18n; +use aish_llm::{Tool, ToolResult}; + +use super::prompt; + +/// Read file content tool. +pub struct ReadFileTool; + +impl Default for ReadFileTool { + fn default() -> Self { + Self::new() + } +} + +impl ReadFileTool { + pub fn new() -> Self { + Self + } +} + +impl Tool for ReadFileTool { + fn name(&self) -> &str { + "read_file" + } + + fn description(&self) -> &str { + prompt::DESCRIPTION + } + + fn parameters(&self) -> serde_json::Value { + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT + } + + fn execute(&self, args: serde_json::Value) -> ToolResult { + let path = match args.get("path").and_then(|p| p.as_str()) { + Some(p) => p, + None => return ToolResult::error(aish_i18n::t("tools.fs.read_file.missing_path")), + }; + + let raw_bytes = match std::fs::read(path) { + Ok(b) => b, + Err(e) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("error".to_string(), e.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.read_file.read_failed", + &args_map, + )); + } + }; + + const SIZE_LIMIT: usize = 32 * 1024; + if raw_bytes.len() > SIZE_LIMIT { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("size".to_string(), raw_bytes.len().to_string()); + args_map.insert("limit".to_string(), SIZE_LIMIT.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.read_file.file_too_large", + &args_map, + )); + } + + let content = match String::from_utf8(raw_bytes) { + Ok(s) => s, + Err(e) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("error".to_string(), e.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.read_file.decode_failed", + &args_map, + )); + } + }; + + let lines: Vec<&str> = content.lines().collect(); + if lines.is_empty() { + return ToolResult::success(aish_i18n::t("tools.fs.read_file.empty_file")); + } + + let offset = args.get("offset").and_then(|o| o.as_u64()).unwrap_or(0) as usize; + let limit = args + .get("limit") + .and_then(|l| l.as_u64()) + .map(|l| l as usize); + + if offset >= lines.len() { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("offset".to_string(), offset.to_string()); + args_map.insert("length".to_string(), lines.len().to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.read_file.offset_exceeds_length", + &args_map, + )); + } + + let selected: Vec = if let Some(limit) = limit { + lines + .iter() + .skip(offset) + .take(limit) + .enumerate() + .map(|(i, line)| format!("{:>6}\t{}", offset + i + 1, line)) + .collect() + } else { + lines + .iter() + .skip(offset) + .enumerate() + .map(|(i, line)| format!("{:>6}\t{}", offset + i + 1, line)) + .collect() + }; + + ToolResult::success(selected.join("\n")) + } +} + +/// Path-restricted wrapper around [`ReadFileTool`] for SSH sessions. +pub struct SshReadFileTool { + inner: ReadFileTool, + offload_root: std::path::PathBuf, +} + +impl SshReadFileTool { + pub fn new() -> Self { + let offload_root = std::env::temp_dir().join("aish-offload"); + let canonical_root = std::fs::canonicalize(&offload_root).unwrap_or(offload_root); + Self { + inner: ReadFileTool::new(), + offload_root: canonical_root, + } + } +} + +impl Tool for SshReadFileTool { + fn name(&self) -> &str { + self.inner.name() + } + + fn description(&self) -> &str { + self.inner.description() + } + + fn parameters(&self) -> serde_json::Value { + self.inner.parameters() + } + + fn prompt(&self) -> &str { + self.inner.prompt() + } + + fn execute(&self, args: serde_json::Value) -> ToolResult { + let path = match args.get("path").and_then(|p| p.as_str()) { + Some(p) => p, + None => return ToolResult::error(aish_i18n::t("tools.fs.read_file.missing_path")), + }; + let canonical = match std::fs::canonicalize(path) { + Ok(c) => c, + Err(e) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("error".to_string(), e.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.read_file.read_failed", + &args_map, + )); + } + }; + if !canonical.starts_with(&self.offload_root) { + return ToolResult::error("Access denied: path is not inside offload directory"); + } + let mut safe_args = args; + if let Some(obj) = safe_args.as_object_mut() { + obj.insert( + "path".to_string(), + serde_json::Value::String(canonical.to_string_lossy().into_owned()), + ); + } + self.inner.execute(safe_args) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aish_llm::Tool; + use std::fs; + + fn temp_dir() -> tempfile::TempDir { + tempfile::tempdir().expect("failed to create temp dir") + } + + #[test] + fn test_read_file_with_line_numbers() { + let dir = temp_dir(); + let file_path = dir.path().join("test.txt"); + fs::write(&file_path, "hello\nworld\nfoo").unwrap(); + + let tool = ReadFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap() + })); + + assert!(result.ok); + assert_eq!(result.output, " 1\thello\n 2\tworld\n 3\tfoo"); + } + + #[test] + fn test_read_file_with_offset() { + let dir = temp_dir(); + let file_path = dir.path().join("test.txt"); + fs::write(&file_path, "line1\nline2\nline3\nline4\nline5").unwrap(); + + let tool = ReadFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap(), + "offset": 2, + "limit": 2 + })); + + assert!(result.ok); + assert_eq!(result.output, " 3\tline3\n 4\tline4"); + } + + #[test] + fn test_read_file_size_limit() { + aish_i18n::set_locale("en-US"); + + let dir = temp_dir(); + let file_path = dir.path().join("big.txt"); + let big_content = "x".repeat(33 * 1024); + fs::write(&file_path, &big_content).unwrap(); + + let tool = ReadFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap() + })); + + assert!(!result.ok); + assert!( + result.output.contains("limit") || result.output.contains("bytes"), + "Expected size limit error, got: {}", + result.output + ); + } +} diff --git a/crates/aish-tools/src/secure_bash.rs b/crates/aish-tools/src/secure_bash/secure_bash.rs similarity index 98% rename from crates/aish-tools/src/secure_bash.rs rename to crates/aish-tools/src/secure_bash/secure_bash.rs index 67929389..e30755c1 100644 --- a/crates/aish-tools/src/secure_bash.rs +++ b/crates/aish-tools/src/secure_bash/secure_bash.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use aish_llm::{CancellationToken, Tool, ToolResult}; use aish_security::SecurityDecision; -use super::bash::PtySlot; +use crate::bash::PtySlot; /// Type of the security check callback. type SecurityCheckFn = Box SecurityDecision + Send + Sync>; @@ -110,6 +110,10 @@ impl Tool for SecureBashTool { self.inner.parameters() } + fn prompt(&self) -> &str { + self.inner.prompt() + } + fn preflight(&self, args: &serde_json::Value) -> aish_llm::PreflightResult { let command = match args.get("command").and_then(|c| c.as_str()) { Some(cmd) => cmd, diff --git a/crates/aish-tools/src/skill_tool/prompt.rs b/crates/aish-tools/src/skill_tool/prompt.rs new file mode 100644 index 00000000..2084514c --- /dev/null +++ b/crates/aish-tools/src/skill_tool/prompt.rs @@ -0,0 +1,25 @@ +pub(crate) const DESCRIPTION: &str = "Invoke a skill within the main conversation."; + +pub(crate) const PROMPT: &str = r#"Use this tool to invoke user-available skills. + +Usage: +- Invoke a skill before answering when it directly matches the user's request. +- Pass only concise arguments needed by the selected skill. +- Do not invent skill names; use only skills that are available in the current session."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "skill_name": { + "type": "string", + "description": "Name of the skill to invoke." + }, + "args": { + "type": "string", + "description": "Optional arguments for the skill." + } + }, + "required": ["skill_name"] + }) +} diff --git a/crates/aish-tools/src/skill_tool.rs b/crates/aish-tools/src/skill_tool/skill_tool.rs similarity index 86% rename from crates/aish-tools/src/skill_tool.rs rename to crates/aish-tools/src/skill_tool/skill_tool.rs index 69205bf1..f9841805 100644 --- a/crates/aish-tools/src/skill_tool.rs +++ b/crates/aish-tools/src/skill_tool/skill_tool.rs @@ -1,5 +1,7 @@ use aish_llm::{Tool, ToolResult}; +use super::prompt; + /// Callback type for looking up a skill by name. pub type SkillLookupFn = Box Option + Send + Sync>; pub type SkillListFn = Box Vec + Send + Sync>; @@ -39,26 +41,15 @@ impl Tool for SkillTool { } fn description(&self) -> &str { - "Execute a skill within the main conversation. Skills provide specialized capabilities \ - and domain knowledge. When a skill matches the user's request, invoke this tool BEFORE \ - generating any other response." + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "skill_name": { - "type": "string", - "description": "The skill name to invoke. E.g., 'commit', 'review-pr', etc." - }, - "args": { - "type": "string", - "description": "Optional arguments for the skill" - } - }, - "required": ["skill_name"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/system_diagnose/prompt.rs b/crates/aish-tools/src/system_diagnose/prompt.rs new file mode 100644 index 00000000..b803def5 --- /dev/null +++ b/crates/aish-tools/src/system_diagnose/prompt.rs @@ -0,0 +1,21 @@ +pub(crate) const DESCRIPTION: &str = "Run an isolated system diagnosis agent."; + +pub(crate) const PROMPT: &str = r#"Use this tool for deeper system diagnosis tasks. + +Usage: +- Use it when the user asks why a system, service, filesystem, network, or performance issue is happening. +- Provide a concise query describing the symptoms and relevant context. +- Do not use it for simple one-command checks that can be handled directly."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Diagnostic query or system issue to analyze." + } + }, + "required": ["query"] + }) +} diff --git a/crates/aish-tools/src/system_diagnose.rs b/crates/aish-tools/src/system_diagnose/system_diagnose.rs similarity index 92% rename from crates/aish-tools/src/system_diagnose.rs rename to crates/aish-tools/src/system_diagnose/system_diagnose.rs index 6db469c5..9b65c4fb 100644 --- a/crates/aish-tools/src/system_diagnose.rs +++ b/crates/aish-tools/src/system_diagnose/system_diagnose.rs @@ -12,6 +12,8 @@ use aish_llm::diagnose_agent::build_diagnose_prompt; use aish_llm::types::LlmCallbackResult; use aish_llm::{DiagnoseAgent, LlmSession, SubSessionConfig, Tool, ToolResult}; +use super::prompt; + /// Shared event callback holder that can be set after tool construction. /// /// Uses `Mutex>>` because the event callback is created @@ -77,22 +79,15 @@ impl Tool for SystemDiagnoseTool { } fn description(&self) -> &str { - "Advanced log analysis and system diagnosis agent that can read files, \ - analyze patterns, and provide detailed diagnostic reports" + prompt::DESCRIPTION } fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The diagnostic query or system issue to analyze. \ - Describe the problem, symptoms, or specific logs to investigate." - } - }, - "required": ["query"] - }) + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT } fn execute(&self, _args: serde_json::Value) -> ToolResult { diff --git a/crates/aish-tools/src/web_fetch.rs b/crates/aish-tools/src/web_fetch.rs deleted file mode 100644 index 18436b6d..00000000 --- a/crates/aish-tools/src/web_fetch.rs +++ /dev/null @@ -1,935 +0,0 @@ -use std::collections::HashMap; -use std::future::Future; -use std::net::IpAddr; -use std::pin::Pin; -use std::sync::{Mutex, OnceLock}; -use std::time::{Duration, Instant}; - -use aish_llm::{ - ChatMessage, LlmResponse, LlmSession, PreflightResult, PreflightSecurityContext, - SecurityPanelMode, StreamParser, Tool, ToolResult, -}; -use futures::StreamExt; -use regex::Regex; -use reqwest::header::{ACCEPT, CONTENT_TYPE, USER_AGENT}; -use reqwest::{redirect, Client, StatusCode, Url}; - -const TOOL_NAME: &str = "WebFetch"; -const MAX_URL_LENGTH: usize = 2000; -const MAX_HTTP_CONTENT_LENGTH: usize = 10 * 1024 * 1024; -const FETCH_TIMEOUT_SECS: u64 = 60; -const MAX_REDIRECTS: usize = 10; -const MAX_MARKDOWN_LENGTH: usize = 100_000; -const CACHE_TTL: Duration = Duration::from_secs(15 * 60); -const CACHE_MAX_ENTRIES: usize = 64; -const USER_AGENT_VALUE: &str = concat!("aish/", env!("CARGO_PKG_VERSION"), " WebFetch"); - -#[derive(Clone)] -struct CacheEntry { - fetched_at: Instant, - url: String, - code: u16, - code_text: String, - bytes: usize, - content_type: String, - content: String, -} - -static URL_CACHE: OnceLock>> = OnceLock::new(); -static DESCRIPTION: OnceLock = OnceLock::new(); - -fn cache() -> &'static Mutex> { - URL_CACHE.get_or_init(|| Mutex::new(HashMap::new())) -} - -fn get_description() -> &'static str { - DESCRIPTION.get_or_init(|| aish_i18n::t("tools.web_fetch.description")) -} - -/// Fetch a URL, extract readable text, and answer a focused prompt about it. -pub struct WebFetchTool { - api_base: String, - api_key: String, - model: String, - temperature: Option, - max_tokens: Option, -} - -impl WebFetchTool { - pub fn new( - api_base: &str, - api_key: &str, - model: &str, - temperature: Option, - max_tokens: Option, - ) -> Self { - Self { - api_base: api_base.to_string(), - api_key: api_key.to_string(), - model: model.to_string(), - temperature, - max_tokens, - } - } - - fn build_client() -> Result { - Client::builder() - .redirect(redirect::Policy::none()) - .timeout(Duration::from_secs(FETCH_TIMEOUT_SECS)) - .build() - .map_err(|error| error.to_string()) - } - - async fn fetch_url_content(&self, raw_url: &str) -> Result { - let normalized = validate_and_normalize_url(raw_url).map_err(FetchFailure::Blocked)?; - ensure_public_host(&normalized) - .await - .map_err(FetchFailure::Blocked)?; - - if let Some(entry) = get_cached(raw_url) { - return Ok(FetchedContent { - url: entry.url, - code: entry.code, - code_text: entry.code_text, - bytes: entry.bytes, - content_type: entry.content_type, - content: entry.content, - from_cache: true, - }); - } - - let client = Self::build_client().map_err(FetchFailure::Request)?; - let response = get_with_permitted_redirects(&client, normalized.clone(), 0).await?; - let status = response.status(); - let content_type = response - .headers() - .get(CONTENT_TYPE) - .and_then(|value| value.to_str().ok()) - .unwrap_or("") - .to_string(); - - if is_binary_content_type(&content_type) { - return Err(FetchFailure::Request(aish_i18n::t( - "tools.web_fetch.binary_unsupported", - ))); - } - - let raw = read_limited_body(response) - .await - .map_err(FetchFailure::Request)?; - let bytes = raw.len(); - let text = String::from_utf8_lossy(&raw).to_string(); - let content = if content_type.to_ascii_lowercase().contains("text/html") { - html_to_readable_text(&text) - } else { - text - }; - - let entry = CacheEntry { - fetched_at: Instant::now(), - url: normalized.to_string(), - code: status.as_u16(), - code_text: status_text(status).to_string(), - bytes, - content_type: content_type.clone(), - content: content.clone(), - }; - set_cached(raw_url.to_string(), entry); - - Ok(FetchedContent { - url: normalized.to_string(), - code: status.as_u16(), - code_text: status_text(status).to_string(), - bytes, - content_type, - content, - from_cache: false, - }) - } - - async fn apply_prompt_to_content( - &self, - prompt: &str, - content: &str, - is_preapproved_domain: bool, - ) -> Result { - let model_prompt = make_secondary_model_prompt(content, prompt, is_preapproved_domain); - let session = LlmSession::new( - &self.api_base, - &self.api_key, - &self.model, - self.temperature.or(Some(0.1)), - self.max_tokens.or(Some(2048)), - ); - let messages = vec![ChatMessage::system(""), ChatMessage::user(model_prompt)]; - match session - .chat_completion_raw(&messages, None, false, Some(0.1), Some(2048)) - .await - .map_err(|error| error.to_string())? - { - LlmResponse::Json(json) => { - let (content, _reasoning, _tool_calls, _usage) = - StreamParser::parse_response(&json); - Ok(content.unwrap_or_else(|| "No response from model".to_string())) - } - LlmResponse::Stream(_) => Err("secondary model unexpectedly returned a stream".into()), - } - } -} - -impl Tool for WebFetchTool { - fn name(&self) -> &str { - TOOL_NAME - } - - fn description(&self) -> &str { - get_description() - } - - fn parameters(&self) -> serde_json::Value { - serde_json::json!({ - "type": "object", - "properties": { - "url": { - "type": "string", - "description": "The fully-qualified URL to fetch content from" - }, - "prompt": { - "type": "string", - "description": "The prompt describing what information to extract from the fetched page" - } - }, - "required": ["url", "prompt"], - "additionalProperties": false - }) - } - - fn preflight(&self, args: &serde_json::Value) -> PreflightResult { - let url = match args.get("url").and_then(|value| value.as_str()) { - Some(value) if !value.trim().is_empty() => value, - _ => { - return PreflightResult::Block { - message: aish_i18n::t("tools.web_fetch.missing_url"), - security: Some(PreflightSecurityContext::fallback( - TOOL_NAME, - None, - aish_i18n::t("tools.web_fetch.missing_url"), - SecurityPanelMode::Blocked, - )), - } - } - }; - - let normalized = match validate_and_normalize_url(url) { - Ok(parsed) => parsed, - Err(message) => { - return PreflightResult::Block { - message: message.clone(), - security: Some(PreflightSecurityContext::fallback( - TOOL_NAME, - Some(url.to_string()), - message, - SecurityPanelMode::Blocked, - )), - } - } - }; - - let hostname = normalized.host_str().unwrap_or("").to_string(); - if is_preapproved_host(&hostname, normalized.path()) { - return PreflightResult::Allow; - } - - let message = aish_i18n::t_with_args( - "tools.web_fetch.confirm_fetch", - &HashMap::from([("host".to_string(), hostname.clone())]), - ); - PreflightResult::Confirm { - message: message.clone(), - security: Some(PreflightSecurityContext::fallback( - TOOL_NAME, - Some(hostname), - message, - SecurityPanelMode::Confirm, - )), - } - } - - fn execute(&self, _args: serde_json::Value) -> ToolResult { - ToolResult::error("WebFetch requires async execution; use execute_async") - } - - fn execute_async<'a>( - &'a self, - args: serde_json::Value, - ) -> Pin + Send + 'a>> { - Box::pin(async move { - let url = match args.get("url").and_then(|value| value.as_str()) { - Some(value) if !value.trim().is_empty() => value.trim(), - _ => return ToolResult::error(aish_i18n::t("tools.web_fetch.missing_url")), - }; - let prompt = match args.get("prompt").and_then(|value| value.as_str()) { - Some(value) if !value.trim().is_empty() => value.trim(), - _ => return ToolResult::error(aish_i18n::t("tools.web_fetch.missing_prompt")), - }; - - let start = Instant::now(); - let fetched = match self.fetch_url_content(url).await { - Ok(content) => content, - Err(FetchFailure::Redirect(info)) => { - let message = format_redirect_message(&info, prompt); - return ToolResult { - ok: true, - output: message.clone(), - meta: Some(serde_json::json!({ - "url": url, - "redirect_url": info.redirect_url, - "code": info.status_code, - "result": message, - "durationMs": start.elapsed().as_millis() as u64, - })), - }; - } - Err(FetchFailure::Blocked(message)) | Err(FetchFailure::Request(message)) => { - return ToolResult::error(message) - } - }; - - let truncated_content = truncate_for_model(&fetched.content); - let parsed_url = Url::parse(&fetched.url).ok(); - let is_preapproved_domain = parsed_url - .as_ref() - .and_then(|parsed| parsed.host_str().map(|host| (host, parsed.path()))) - .is_some_and(|(host, path)| is_preapproved_host(host, path)); - - let result = match self - .apply_prompt_to_content(prompt, &truncated_content, is_preapproved_domain) - .await - { - Ok(result) => result, - Err(error) => { - let mut args_map = HashMap::new(); - args_map.insert("error".to_string(), error); - return ToolResult::error(aish_i18n::t_with_args( - "tools.web_fetch.secondary_model_failed", - &args_map, - )); - } - }; - - let duration_ms = start.elapsed().as_millis() as u64; - let output = format!( - "Fetched: {}\nStatus: {} {}\nBytes: {}\nDuration: {}ms\nCached: {}\n\n{}", - fetched.url, - fetched.code, - fetched.code_text, - fetched.bytes, - duration_ms, - fetched.from_cache, - result - ); - - ToolResult { - ok: true, - output, - meta: Some(serde_json::json!({ - "url": fetched.url, - "code": fetched.code, - "codeText": fetched.code_text, - "bytes": fetched.bytes, - "contentType": fetched.content_type, - "durationMs": duration_ms, - "fromCache": fetched.from_cache, - "result": result, - })), - } - }) - } -} - -#[derive(Debug)] -enum FetchFailure { - Blocked(String), - Request(String), - Redirect(RedirectInfo), -} - -#[derive(Debug)] -struct RedirectInfo { - original_url: String, - redirect_url: String, - status_code: u16, -} - -struct FetchedContent { - url: String, - code: u16, - code_text: String, - bytes: usize, - content_type: String, - content: String, - from_cache: bool, -} - -async fn get_with_permitted_redirects( - client: &Client, - url: Url, - depth: usize, -) -> Result { - if depth > MAX_REDIRECTS { - return Err(FetchFailure::Request(format!( - "Too many redirects (exceeded {})", - MAX_REDIRECTS - ))); - } - - ensure_public_host(&url) - .await - .map_err(FetchFailure::Blocked)?; - let response = client - .get(url.clone()) - .header(ACCEPT, "text/markdown, text/html, text/plain, */*") - .header(USER_AGENT, USER_AGENT_VALUE) - .send() - .await - .map_err(|error| FetchFailure::Request(error.to_string()))?; - - if is_redirect_status(response.status()) { - let location = response - .headers() - .get(reqwest::header::LOCATION) - .and_then(|value| value.to_str().ok()) - .ok_or_else(|| FetchFailure::Request("Redirect missing Location header".into()))?; - let redirect_url = url - .join(location) - .map_err(|error| FetchFailure::Request(error.to_string()))?; - validate_url_basics(&redirect_url).map_err(FetchFailure::Blocked)?; - ensure_public_host(&redirect_url) - .await - .map_err(FetchFailure::Blocked)?; - - if is_permitted_redirect(&url, &redirect_url) { - return Box::pin(get_with_permitted_redirects( - client, - redirect_url, - depth + 1, - )) - .await; - } - - return Err(FetchFailure::Redirect(RedirectInfo { - original_url: url.to_string(), - redirect_url: redirect_url.to_string(), - status_code: response.status().as_u16(), - })); - } - - Ok(response) -} - -fn validate_and_normalize_url(raw_url: &str) -> Result { - if raw_url.len() > MAX_URL_LENGTH { - return Err(aish_i18n::t("tools.web_fetch.invalid_url")); - } - let mut parsed = - Url::parse(raw_url).map_err(|_| aish_i18n::t("tools.web_fetch.invalid_url"))?; - if parsed.scheme() == "http" { - parsed - .set_scheme("https") - .map_err(|_| aish_i18n::t("tools.web_fetch.invalid_url"))?; - } - validate_url_basics(&parsed)?; - Ok(parsed) -} - -fn validate_url_basics(parsed: &Url) -> Result<(), String> { - if parsed.scheme() != "https" && parsed.scheme() != "http" { - return Err(aish_i18n::t("tools.web_fetch.invalid_url")); - } - if !parsed.username().is_empty() || parsed.password().is_some() { - return Err(aish_i18n::t("tools.web_fetch.invalid_url")); - } - let host = parsed - .host_str() - .ok_or_else(|| aish_i18n::t("tools.web_fetch.invalid_url"))?; - if host.split('.').count() < 2 && host.parse::().is_err() { - return Err(aish_i18n::t("tools.web_fetch.invalid_url")); - } - if is_blocked_hostname(host) { - let mut args = HashMap::new(); - args.insert("host".to_string(), host.to_string()); - return Err(aish_i18n::t_with_args( - "tools.web_fetch.blocked_private_host", - &args, - )); - } - Ok(()) -} - -async fn ensure_public_host(url: &Url) -> Result<(), String> { - let host = url - .host_str() - .ok_or_else(|| aish_i18n::t("tools.web_fetch.invalid_url"))?; - if is_blocked_hostname(host) { - let mut args = HashMap::new(); - args.insert("host".to_string(), host.to_string()); - return Err(aish_i18n::t_with_args( - "tools.web_fetch.blocked_private_host", - &args, - )); - } - if host.parse::().is_ok() { - return Ok(()); - } - - let port = url.port_or_known_default().unwrap_or(443); - let addrs = tokio::net::lookup_host((host, port)) - .await - .map_err(|error| format!("DNS lookup failed for {}: {}", host, error))?; - let mut found = false; - for addr in addrs { - found = true; - if is_private_ip(&addr.ip()) { - let mut args = HashMap::new(); - args.insert("host".to_string(), host.to_string()); - return Err(aish_i18n::t_with_args( - "tools.web_fetch.blocked_private_host", - &args, - )); - } - } - if !found { - return Err(format!("DNS lookup returned no addresses for {}", host)); - } - Ok(()) -} - -fn is_blocked_hostname(host: &str) -> bool { - let normalized = host.trim_end_matches('.').to_ascii_lowercase(); - if matches!( - normalized.as_str(), - "localhost" | "metadata.google.internal" - ) { - return true; - } - if normalized.ends_with(".localhost") || normalized.ends_with(".local") { - return true; - } - match normalized.parse::() { - Ok(ip) => is_private_ip(&ip), - Err(_) => false, - } -} - -fn is_private_ip(ip: &IpAddr) -> bool { - match ip { - IpAddr::V4(addr) => { - addr.is_private() - || addr.is_loopback() - || addr.is_link_local() - || addr.is_broadcast() - || addr.is_documentation() - || addr.is_unspecified() - || addr.octets() == [169, 254, 169, 254] - || addr.octets()[0] == 0 - } - IpAddr::V6(addr) => { - addr.is_loopback() - || addr.is_unspecified() - || addr.is_unique_local() - || addr.is_unicast_link_local() - } - } -} - -fn is_redirect_status(status: StatusCode) -> bool { - matches!( - status, - StatusCode::MOVED_PERMANENTLY - | StatusCode::FOUND - | StatusCode::TEMPORARY_REDIRECT - | StatusCode::PERMANENT_REDIRECT - ) -} - -fn is_permitted_redirect(original: &Url, redirect_url: &Url) -> bool { - if original.scheme() != redirect_url.scheme() || original.port() != redirect_url.port() { - return false; - } - if !redirect_url.username().is_empty() || redirect_url.password().is_some() { - return false; - } - let Some(original_host) = original.host_str() else { - return false; - }; - let Some(redirect_host) = redirect_url.host_str() else { - return false; - }; - strip_www(original_host) == strip_www(redirect_host) -} - -fn strip_www(host: &str) -> &str { - host.strip_prefix("www.").unwrap_or(host) -} - -async fn read_limited_body(response: reqwest::Response) -> Result, String> { - if response - .content_length() - .is_some_and(|length| length > MAX_HTTP_CONTENT_LENGTH as u64) - { - return Err(aish_i18n::t("tools.web_fetch.content_too_large")); - } - - let mut stream = response.bytes_stream(); - let mut buffer = Vec::new(); - while let Some(chunk) = stream.next().await { - let chunk = chunk.map_err(|error| error.to_string())?; - if buffer.len() + chunk.len() > MAX_HTTP_CONTENT_LENGTH { - return Err(aish_i18n::t("tools.web_fetch.content_too_large")); - } - buffer.extend_from_slice(&chunk); - } - Ok(buffer) -} - -fn status_text(status: StatusCode) -> &'static str { - status.canonical_reason().unwrap_or("") -} - -fn is_binary_content_type(content_type: &str) -> bool { - let lower = content_type.to_ascii_lowercase(); - if lower.starts_with("text/") { - return false; - } - if lower.contains("json") || lower.contains("xml") || lower.contains("javascript") { - return false; - } - lower.contains("application/pdf") - || lower.starts_with("image/") - || lower.starts_with("audio/") - || lower.starts_with("video/") - || lower.contains("application/octet-stream") -} - -fn html_to_readable_text(html: &str) -> String { - let without_scripts = regex_replace_all( - html, - r"(?is)]*>.*?|]*>.*?|]*>.*?|]*>.*?|]*>.*?", - "\n", - ); - let with_breaks = regex_replace_all( - &without_scripts, - r"(?i)]*>", - "\n", - ); - let without_tags = regex_replace_all(&with_breaks, r"(?is)<[^>]+>", " "); - normalize_text_whitespace(&decode_html_entities(&without_tags)) -} - -fn regex_replace_all(input: &str, pattern: &str, replacement: &str) -> String { - match Regex::new(pattern) { - Ok(regex) => regex.replace_all(input, replacement).to_string(), - Err(_) => input.to_string(), - } -} - -fn decode_html_entities(input: &str) -> String { - input - .replace(" ", " ") - .replace("&", "&") - .replace("<", "<") - .replace(">", ">") - .replace(""", "\"") - .replace("'", "'") - .replace("'", "'") -} - -fn normalize_text_whitespace(input: &str) -> String { - let mut lines = Vec::new(); - for line in input.lines() { - let collapsed = line.split_whitespace().collect::>().join(" "); - if !collapsed.is_empty() { - lines.push(collapsed); - } - } - lines.join("\n") -} - -fn truncate_for_model(content: &str) -> String { - if content.chars().count() <= MAX_MARKDOWN_LENGTH { - return content.to_string(); - } - let mut truncated = content - .chars() - .take(MAX_MARKDOWN_LENGTH) - .collect::(); - truncated.push_str("\n\n[Content truncated due to length...]"); - truncated -} - -fn make_secondary_model_prompt( - markdown_content: &str, - prompt: &str, - is_preapproved_domain: bool, -) -> String { - let guidelines = if is_preapproved_domain { - "Provide a concise response based on the content above. Include relevant details, code examples, and documentation excerpts as needed." - } else { - "Provide a concise response based only on the content above. In your response:\n - Enforce a strict 125-character maximum for quotes from any source document. Open Source Software is ok as long as we respect the license.\n - Use quotation marks for exact language from articles; any language outside of the quotation should never be word-for-word the same.\n - You are not a lawyer and never comment on the legality of your own prompts and responses.\n - Never produce or reproduce exact song lyrics." - }; - format!( - "Web page content:\n---\n{}\n---\n\n{}\n\n{}\n", - markdown_content, prompt, guidelines - ) -} - -fn format_redirect_message(info: &RedirectInfo, prompt: &str) -> String { - let status_text = match info.status_code { - 301 => "Moved Permanently", - 307 => "Temporary Redirect", - 308 => "Permanent Redirect", - _ => "Found", - }; - format!( - "REDIRECT DETECTED: The URL redirects to a different host.\n\nOriginal URL: {}\nRedirect URL: {}\nStatus: {} {}\n\nTo complete your request, fetch the redirected URL with these parameters:\n- url: \"{}\"\n- prompt: \"{}\"", - info.original_url, - info.redirect_url, - info.status_code, - status_text, - info.redirect_url, - prompt - ) -} - -fn get_cached(key: &str) -> Option { - let mut guard = cache().lock().ok()?; - let now = Instant::now(); - guard.retain(|_, entry| now.duration_since(entry.fetched_at) < CACHE_TTL); - guard.get(key).cloned() -} - -fn set_cached(key: String, entry: CacheEntry) { - if let Ok(mut guard) = cache().lock() { - if guard.len() >= CACHE_MAX_ENTRIES { - if let Some(oldest_key) = guard - .iter() - .min_by_key(|(_, value)| value.fetched_at) - .map(|(cache_key, _)| cache_key.clone()) - { - guard.remove(&oldest_key); - } - } - guard.insert(key, entry); - } -} - -fn is_preapproved_host(hostname: &str, pathname: &str) -> bool { - for entry in PREAPPROVED_HOSTS { - if let Some((host, prefix)) = entry.split_once('/') { - if hostname == host { - let prefix = format!("/{}", prefix); - if pathname == prefix || pathname.starts_with(&(prefix + "/")) { - return true; - } - } - continue; - } - if hostname == *entry { - return true; - } - } - false -} - -const PREAPPROVED_HOSTS: &[&str] = &[ - "platform.claude.com", - "code.claude.com", - "modelcontextprotocol.io", - "github.com/anthropics", - "agentskills.io", - "docs.python.org", - "en.cppreference.com", - "docs.oracle.com", - "learn.microsoft.com", - "developer.mozilla.org", - "go.dev", - "pkg.go.dev", - "www.php.net", - "docs.swift.org", - "kotlinlang.org", - "ruby-doc.org", - "doc.rust-lang.org", - "www.typescriptlang.org", - "react.dev", - "angular.io", - "vuejs.org", - "nextjs.org", - "expressjs.com", - "nodejs.org", - "bun.sh", - "jquery.com", - "getbootstrap.com", - "tailwindcss.com", - "d3js.org", - "threejs.org", - "redux.js.org", - "webpack.js.org", - "jestjs.io", - "reactrouter.com", - "docs.djangoproject.com", - "flask.palletsprojects.com", - "fastapi.tiangolo.com", - "pandas.pydata.org", - "numpy.org", - "www.tensorflow.org", - "pytorch.org", - "scikit-learn.org", - "matplotlib.org", - "requests.readthedocs.io", - "jupyter.org", - "laravel.com", - "symfony.com", - "wordpress.org", - "docs.spring.io", - "hibernate.org", - "tomcat.apache.org", - "gradle.org", - "maven.apache.org", - "asp.net", - "dotnet.microsoft.com", - "nuget.org", - "blazor.net", - "reactnative.dev", - "docs.flutter.dev", - "developer.apple.com", - "developer.android.com", - "keras.io", - "spark.apache.org", - "huggingface.co", - "www.kaggle.com", - "www.mongodb.com", - "redis.io", - "www.postgresql.org", - "dev.mysql.com", - "www.sqlite.org", - "graphql.org", - "prisma.io", - "docs.aws.amazon.com", - "cloud.google.com", - "kubernetes.io", - "www.docker.com", - "www.terraform.io", - "www.ansible.com", - "vercel.com/docs", - "docs.netlify.com", - "devcenter.heroku.com", - "cypress.io", - "selenium.dev", - "docs.unity.com", - "docs.unrealengine.com", - "git-scm.com", - "nginx.org", - "httpd.apache.org", -]; - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn normalizes_http_to_https() { - let url = validate_and_normalize_url("http://example.com/path").unwrap(); - assert_eq!(url.as_str(), "https://example.com/path"); - } - - #[test] - fn rejects_private_hosts() { - assert!(validate_and_normalize_url("https://localhost/").is_err()); - assert!(validate_and_normalize_url("https://127.0.0.1/").is_err()); - assert!(validate_and_normalize_url("https://169.254.169.254/").is_err()); - assert!(validate_and_normalize_url("https://10.0.0.5/").is_err()); - } - - #[test] - fn preapproved_host_supports_path_prefix_boundary() { - assert!(is_preapproved_host("github.com", "/anthropics/claude-code")); - assert!(!is_preapproved_host( - "github.com", - "/anthropics-evil/project" - )); - assert!(is_preapproved_host("doc.rust-lang.org", "/book/")); - } - - #[test] - fn redirect_only_allows_same_origin_or_www_equivalent() { - let original = Url::parse("https://example.com/docs").unwrap(); - let same = Url::parse("https://www.example.com/docs").unwrap(); - let other = Url::parse("https://evil.example.net/docs").unwrap(); - let http = Url::parse("http://example.com/docs").unwrap(); - assert!(is_permitted_redirect(&original, &same)); - assert!(!is_permitted_redirect(&original, &other)); - assert!(!is_permitted_redirect(&original, &http)); - } - - #[test] - fn html_to_readable_text_removes_scripts_and_tags() { - let html = "

Hello & hi

World

"; - let text = html_to_readable_text(html); - assert!(text.contains("Hello & hi")); - assert!(text.contains("World")); - assert!(!text.contains("bad()")); - assert!(!text.contains("

")); - } - - #[test] - fn secondary_prompt_includes_quote_restriction_for_unapproved_domains() { - let prompt = make_secondary_model_prompt("content", "summarize", false); - assert!(prompt.contains("125-character maximum")); - assert!(prompt.contains("summarize")); - } - - #[tokio::test] - #[ignore] - async fn live_fetch_url() { - if std::env::var("AISH_LIVE_WEBFETCH").ok().as_deref() != Some("1") { - eprintln!("set AISH_LIVE_WEBFETCH=1 to run this live network smoke test"); - return; - } - - let url = std::env::var("AISH_LIVE_WEBFETCH_URL") - .unwrap_or_else(|_| "https://github.com/mattpocock/skills".to_string()); - let expected = std::env::var("AISH_LIVE_WEBFETCH_EXPECT").ok(); - let tool = WebFetchTool::new("", "", "", Some(0.1), Some(256)); - let fetched = tool - .fetch_url_content(&url) - .await - .expect("expected live page fetch to succeed"); - - println!( - "fetched {} status={} bytes={} content_type={} chars={}", - fetched.url, - fetched.code, - fetched.bytes, - fetched.content_type, - fetched.content.len() - ); - println!( - "preview:\n{}", - truncate_for_model(&fetched.content) - .chars() - .take(800) - .collect::() - ); - - assert_eq!(fetched.code, 200); - if let Some(expected) = expected { - assert!( - fetched - .content - .to_ascii_lowercase() - .contains(&expected.to_ascii_lowercase()), - "expected fetched content to contain {expected:?}" - ); - } - } -} diff --git a/crates/aish-tools/src/web_fetch/preapproved.rs b/crates/aish-tools/src/web_fetch/preapproved.rs new file mode 100644 index 00000000..2493b8d1 --- /dev/null +++ b/crates/aish-tools/src/web_fetch/preapproved.rs @@ -0,0 +1,127 @@ +/// Domains that can be fetched without an extra WebFetch confirmation. +/// +/// This list is intentionally scoped to WebFetch GET requests. It should not be +/// reused as a general network allowlist for shell commands or sandbox egress. +const PREAPPROVED_HOSTS: &[&str] = &[ + "platform.claude.com", + "code.claude.com", + "modelcontextprotocol.io", + "github.com/anthropics", + "agentskills.io", + "docs.python.org", + "en.cppreference.com", + "docs.oracle.com", + "learn.microsoft.com", + "developer.mozilla.org", + "go.dev", + "pkg.go.dev", + "www.php.net", + "docs.swift.org", + "kotlinlang.org", + "ruby-doc.org", + "doc.rust-lang.org", + "www.typescriptlang.org", + "react.dev", + "angular.io", + "vuejs.org", + "nextjs.org", + "expressjs.com", + "nodejs.org", + "bun.sh", + "jquery.com", + "getbootstrap.com", + "tailwindcss.com", + "d3js.org", + "threejs.org", + "redux.js.org", + "webpack.js.org", + "jestjs.io", + "reactrouter.com", + "docs.djangoproject.com", + "flask.palletsprojects.com", + "fastapi.tiangolo.com", + "pandas.pydata.org", + "numpy.org", + "www.tensorflow.org", + "pytorch.org", + "scikit-learn.org", + "matplotlib.org", + "requests.readthedocs.io", + "jupyter.org", + "laravel.com", + "symfony.com", + "wordpress.org", + "docs.spring.io", + "hibernate.org", + "tomcat.apache.org", + "gradle.org", + "maven.apache.org", + "asp.net", + "dotnet.microsoft.com", + "nuget.org", + "blazor.net", + "reactnative.dev", + "docs.flutter.dev", + "developer.apple.com", + "developer.android.com", + "keras.io", + "spark.apache.org", + "huggingface.co", + "www.kaggle.com", + "www.mongodb.com", + "redis.io", + "www.postgresql.org", + "dev.mysql.com", + "www.sqlite.org", + "graphql.org", + "prisma.io", + "docs.aws.amazon.com", + "cloud.google.com", + "kubernetes.io", + "www.docker.com", + "www.terraform.io", + "www.ansible.com", + "vercel.com/docs", + "docs.netlify.com", + "devcenter.heroku.com", + "cypress.io", + "selenium.dev", + "docs.unity.com", + "docs.unrealengine.com", + "git-scm.com", + "nginx.org", + "httpd.apache.org", +]; + +pub(crate) fn is_preapproved_host(hostname: &str, pathname: &str) -> bool { + for entry in PREAPPROVED_HOSTS { + if let Some((host, prefix)) = entry.split_once('/') { + if hostname == host { + let prefix = format!("/{}", prefix); + if pathname == prefix || pathname.starts_with(&(prefix + "/")) { + return true; + } + } + continue; + } + if hostname == *entry { + return true; + } + } + false +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn preapproved_host_supports_path_prefix_boundary() { + assert!(is_preapproved_host("github.com", "/anthropics/claude-code")); + assert!(!is_preapproved_host( + "github.com", + "/anthropics-evil/project" + )); + assert!(is_preapproved_host("doc.rust-lang.org", "/book/")); + } +} diff --git a/crates/aish-tools/src/web_fetch/prompt.rs b/crates/aish-tools/src/web_fetch/prompt.rs new file mode 100644 index 00000000..1ac7d655 --- /dev/null +++ b/crates/aish-tools/src/web_fetch/prompt.rs @@ -0,0 +1,34 @@ +pub(crate) const DESCRIPTION: &str = + "Fetch public web content and answer a focused prompt about the page."; + +pub(crate) const PROMPT: &str = r#"Use this tool to fetch public web content and answer a focused question about it. + +Usage: +- Use this tool when you need to retrieve and analyze public web content. +- If an authenticated MCP web fetch tool is available, prefer that tool for private or authenticated services. +- WebFetch will fail for authenticated or private URLs such as Google Docs, Confluence, Jira, private GitHub pages, localhost, and internal services. +- Provide a fully-qualified URL and a focused prompt describing what to extract from the page. +- HTTP URLs are automatically upgraded to HTTPS. +- HTML content is converted to readable text before being processed by a secondary model. +- Results may be summarized or truncated when the page is very large. +- A 15-minute cache is used for repeated requests to the same URL. +- If a URL redirects to a different host, call WebFetch again with the redirect URL returned by the tool. +- For GitHub URLs, prefer gh via bash when repository metadata, issues, PRs, releases, or API data are needed."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "url": { + "type": "string", + "description": "Fully-qualified URL to fetch content from." + }, + "prompt": { + "type": "string", + "description": "Focused prompt describing what information to extract from the fetched page." + } + }, + "required": ["url", "prompt"], + "additionalProperties": false + }) +} diff --git a/crates/aish-tools/src/web_fetch/utils.rs b/crates/aish-tools/src/web_fetch/utils.rs new file mode 100644 index 00000000..c0e61bd7 --- /dev/null +++ b/crates/aish-tools/src/web_fetch/utils.rs @@ -0,0 +1,533 @@ +use std::collections::HashMap; +use std::net::IpAddr; +use std::sync::{Mutex, OnceLock}; +use std::time::{Duration, Instant}; + +use futures::StreamExt; +use regex::Regex; +use reqwest::header::{ACCEPT, CONTENT_TYPE, USER_AGENT}; +use reqwest::{redirect, Client, StatusCode, Url}; + +const MAX_URL_LENGTH: usize = 2000; +const MAX_HTTP_CONTENT_LENGTH: usize = 10 * 1024 * 1024; +const FETCH_TIMEOUT_SECS: u64 = 60; +const MAX_REDIRECTS: usize = 10; +pub(crate) const MAX_MARKDOWN_LENGTH: usize = 100_000; +const CACHE_TTL: Duration = Duration::from_secs(15 * 60); +const CACHE_MAX_ENTRIES: usize = 64; +const USER_AGENT_VALUE: &str = concat!("aish/", env!("CARGO_PKG_VERSION"), " WebFetch"); + +#[derive(Clone)] +struct CacheEntry { + fetched_at: Instant, + url: String, + code: u16, + code_text: String, + bytes: usize, + content_type: String, + content: String, +} + +static URL_CACHE: OnceLock>> = OnceLock::new(); + +fn cache() -> &'static Mutex> { + URL_CACHE.get_or_init(|| Mutex::new(HashMap::new())) +} + +#[derive(Debug)] +pub(crate) enum FetchFailure { + Blocked(String), + Request(String), + Redirect(RedirectInfo), +} + +#[derive(Debug)] +pub(crate) struct RedirectInfo { + pub(crate) original_url: String, + pub(crate) redirect_url: String, + pub(crate) status_code: u16, +} + +pub(crate) struct FetchedContent { + pub(crate) url: String, + pub(crate) code: u16, + pub(crate) code_text: String, + pub(crate) bytes: usize, + pub(crate) content_type: String, + pub(crate) content: String, + pub(crate) from_cache: bool, +} + +pub(crate) fn build_client() -> Result { + Client::builder() + .redirect(redirect::Policy::none()) + .timeout(Duration::from_secs(FETCH_TIMEOUT_SECS)) + .build() + .map_err(|error| error.to_string()) +} + +pub(crate) async fn fetch_url_content(raw_url: &str) -> Result { + let normalized = validate_and_normalize_url(raw_url).map_err(FetchFailure::Blocked)?; + ensure_public_host(&normalized) + .await + .map_err(FetchFailure::Blocked)?; + + if let Some(entry) = get_cached(raw_url) { + return Ok(FetchedContent { + url: entry.url, + code: entry.code, + code_text: entry.code_text, + bytes: entry.bytes, + content_type: entry.content_type, + content: entry.content, + from_cache: true, + }); + } + + let client = build_client().map_err(FetchFailure::Request)?; + let response = get_with_permitted_redirects(&client, normalized.clone(), 0).await?; + let status = response.status(); + let content_type = response + .headers() + .get(CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or("") + .to_string(); + + if is_binary_content_type(&content_type) { + return Err(FetchFailure::Request(aish_i18n::t( + "tools.web_fetch.binary_unsupported", + ))); + } + + let raw = read_limited_body(response) + .await + .map_err(FetchFailure::Request)?; + let bytes = raw.len(); + let text = String::from_utf8_lossy(&raw).to_string(); + let content = if content_type.to_ascii_lowercase().contains("text/html") { + html_to_readable_text(&text) + } else { + text + }; + + let entry = CacheEntry { + fetched_at: Instant::now(), + url: normalized.to_string(), + code: status.as_u16(), + code_text: status_text(status).to_string(), + bytes, + content_type: content_type.clone(), + content: content.clone(), + }; + set_cached(raw_url.to_string(), entry); + + Ok(FetchedContent { + url: normalized.to_string(), + code: status.as_u16(), + code_text: status_text(status).to_string(), + bytes, + content_type, + content, + from_cache: false, + }) +} + +async fn get_with_permitted_redirects( + client: &Client, + url: Url, + depth: usize, +) -> Result { + if depth > MAX_REDIRECTS { + return Err(FetchFailure::Request(format!( + "Too many redirects (exceeded {})", + MAX_REDIRECTS + ))); + } + + ensure_public_host(&url) + .await + .map_err(FetchFailure::Blocked)?; + let response = client + .get(url.clone()) + .header(ACCEPT, "text/markdown, text/html, text/plain, */*") + .header(USER_AGENT, USER_AGENT_VALUE) + .send() + .await + .map_err(|error| FetchFailure::Request(error.to_string()))?; + + if is_redirect_status(response.status()) { + let location = response + .headers() + .get(reqwest::header::LOCATION) + .and_then(|value| value.to_str().ok()) + .ok_or_else(|| FetchFailure::Request("Redirect missing Location header".into()))?; + let redirect_url = url + .join(location) + .map_err(|error| FetchFailure::Request(error.to_string()))?; + validate_url_basics(&redirect_url).map_err(FetchFailure::Blocked)?; + ensure_public_host(&redirect_url) + .await + .map_err(FetchFailure::Blocked)?; + + if is_permitted_redirect(&url, &redirect_url) { + return Box::pin(get_with_permitted_redirects( + client, + redirect_url, + depth + 1, + )) + .await; + } + + return Err(FetchFailure::Redirect(RedirectInfo { + original_url: url.to_string(), + redirect_url: redirect_url.to_string(), + status_code: response.status().as_u16(), + })); + } + + Ok(response) +} + +pub(crate) fn validate_and_normalize_url(raw_url: &str) -> Result { + if raw_url.len() > MAX_URL_LENGTH { + return Err(aish_i18n::t("tools.web_fetch.invalid_url")); + } + let mut parsed = + Url::parse(raw_url).map_err(|_| aish_i18n::t("tools.web_fetch.invalid_url"))?; + if parsed.scheme() == "http" { + parsed + .set_scheme("https") + .map_err(|_| aish_i18n::t("tools.web_fetch.invalid_url"))?; + } + validate_url_basics(&parsed)?; + Ok(parsed) +} + +fn validate_url_basics(parsed: &Url) -> Result<(), String> { + if parsed.scheme() != "https" && parsed.scheme() != "http" { + return Err(aish_i18n::t("tools.web_fetch.invalid_url")); + } + if !parsed.username().is_empty() || parsed.password().is_some() { + return Err(aish_i18n::t("tools.web_fetch.invalid_url")); + } + let host = parsed + .host_str() + .ok_or_else(|| aish_i18n::t("tools.web_fetch.invalid_url"))?; + if host.split('.').count() < 2 && host.parse::().is_err() { + return Err(aish_i18n::t("tools.web_fetch.invalid_url")); + } + if is_blocked_hostname(host) { + let mut args = HashMap::new(); + args.insert("host".to_string(), host.to_string()); + return Err(aish_i18n::t_with_args( + "tools.web_fetch.blocked_private_host", + &args, + )); + } + Ok(()) +} + +async fn ensure_public_host(url: &Url) -> Result<(), String> { + let host = url + .host_str() + .ok_or_else(|| aish_i18n::t("tools.web_fetch.invalid_url"))?; + if is_blocked_hostname(host) { + let mut args = HashMap::new(); + args.insert("host".to_string(), host.to_string()); + return Err(aish_i18n::t_with_args( + "tools.web_fetch.blocked_private_host", + &args, + )); + } + if host.parse::().is_ok() { + return Ok(()); + } + + let port = url.port_or_known_default().unwrap_or(443); + let addrs = tokio::net::lookup_host((host, port)) + .await + .map_err(|error| format!("DNS lookup failed for {}: {}", host, error))?; + let mut found = false; + for addr in addrs { + found = true; + if is_private_ip(&addr.ip()) { + let mut args = HashMap::new(); + args.insert("host".to_string(), host.to_string()); + return Err(aish_i18n::t_with_args( + "tools.web_fetch.blocked_private_host", + &args, + )); + } + } + if !found { + return Err(format!("DNS lookup returned no addresses for {}", host)); + } + Ok(()) +} + +fn is_blocked_hostname(host: &str) -> bool { + let normalized = host.trim_end_matches('.').to_ascii_lowercase(); + if matches!( + normalized.as_str(), + "localhost" | "metadata.google.internal" + ) { + return true; + } + if normalized.ends_with(".localhost") || normalized.ends_with(".local") { + return true; + } + match normalized.parse::() { + Ok(ip) => is_private_ip(&ip), + Err(_) => false, + } +} + +fn is_private_ip(ip: &IpAddr) -> bool { + match ip { + IpAddr::V4(addr) => { + addr.is_private() + || addr.is_loopback() + || addr.is_link_local() + || addr.is_broadcast() + || addr.is_documentation() + || addr.is_unspecified() + || addr.octets() == [169, 254, 169, 254] + || addr.octets()[0] == 0 + } + IpAddr::V6(addr) => { + addr.is_loopback() + || addr.is_unspecified() + || addr.is_unique_local() + || addr.is_unicast_link_local() + } + } +} + +fn is_redirect_status(status: StatusCode) -> bool { + matches!( + status, + StatusCode::MOVED_PERMANENTLY + | StatusCode::FOUND + | StatusCode::TEMPORARY_REDIRECT + | StatusCode::PERMANENT_REDIRECT + ) +} + +fn is_permitted_redirect(original: &Url, redirect_url: &Url) -> bool { + if original.scheme() != redirect_url.scheme() || original.port() != redirect_url.port() { + return false; + } + if !redirect_url.username().is_empty() || redirect_url.password().is_some() { + return false; + } + let Some(original_host) = original.host_str() else { + return false; + }; + let Some(redirect_host) = redirect_url.host_str() else { + return false; + }; + strip_www(original_host) == strip_www(redirect_host) +} + +fn strip_www(host: &str) -> &str { + host.strip_prefix("www.").unwrap_or(host) +} + +async fn read_limited_body(response: reqwest::Response) -> Result, String> { + if response + .content_length() + .is_some_and(|length| length > MAX_HTTP_CONTENT_LENGTH as u64) + { + return Err(aish_i18n::t("tools.web_fetch.content_too_large")); + } + + let mut stream = response.bytes_stream(); + let mut buffer = Vec::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|error| error.to_string())?; + if buffer.len() + chunk.len() > MAX_HTTP_CONTENT_LENGTH { + return Err(aish_i18n::t("tools.web_fetch.content_too_large")); + } + buffer.extend_from_slice(&chunk); + } + Ok(buffer) +} + +fn status_text(status: StatusCode) -> &'static str { + status.canonical_reason().unwrap_or("") +} + +fn is_binary_content_type(content_type: &str) -> bool { + let lower = content_type.to_ascii_lowercase(); + if lower.starts_with("text/") { + return false; + } + if lower.contains("json") || lower.contains("xml") || lower.contains("javascript") { + return false; + } + lower.contains("application/pdf") + || lower.starts_with("image/") + || lower.starts_with("audio/") + || lower.starts_with("video/") + || lower.contains("application/octet-stream") +} + +fn html_to_readable_text(html: &str) -> String { + let without_scripts = regex_replace_all( + html, + r"(?is)]*>.*?|]*>.*?|]*>.*?|]*>.*?|]*>.*?", + "\n", + ); + let with_breaks = regex_replace_all( + &without_scripts, + r"(?i)]*>", + "\n", + ); + let without_tags = regex_replace_all(&with_breaks, r"(?is)<[^>]+>", " "); + normalize_text_whitespace(&decode_html_entities(&without_tags)) +} + +fn regex_replace_all(input: &str, pattern: &str, replacement: &str) -> String { + match Regex::new(pattern) { + Ok(regex) => regex.replace_all(input, replacement).to_string(), + Err(_) => input.to_string(), + } +} + +fn decode_html_entities(input: &str) -> String { + input + .replace(" ", " ") + .replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace(""", "\"") + .replace("'", "'") + .replace("'", "'") +} + +fn normalize_text_whitespace(input: &str) -> String { + let mut lines = Vec::new(); + for line in input.lines() { + let collapsed = line.split_whitespace().collect::>().join(" "); + if !collapsed.is_empty() { + lines.push(collapsed); + } + } + lines.join("\n") +} + +pub(crate) fn truncate_for_model(content: &str) -> String { + if content.chars().count() <= MAX_MARKDOWN_LENGTH { + return content.to_string(); + } + let mut truncated = content + .chars() + .take(MAX_MARKDOWN_LENGTH) + .collect::(); + truncated.push_str("\n\n[Content truncated due to length...]"); + truncated +} + +pub(crate) fn format_redirect_message(info: &RedirectInfo, prompt: &str) -> String { + let status_text = match info.status_code { + 301 => "Moved Permanently", + 307 => "Temporary Redirect", + 308 => "Permanent Redirect", + _ => "Found", + }; + format!( + "REDIRECT DETECTED: The URL redirects to a different host.\n\nOriginal URL: {}\nRedirect URL: {}\nStatus: {} {}\n\nTo complete your request, fetch the redirected URL with these parameters:\n- url: \"{}\"\n- prompt: \"{}\"", + info.original_url, + info.redirect_url, + info.status_code, + status_text, + info.redirect_url, + prompt + ) +} + +pub(crate) fn make_secondary_model_prompt( + markdown_content: &str, + prompt: &str, + is_preapproved_domain: bool, +) -> String { + let guidelines = if is_preapproved_domain { + "Provide a concise response based on the content above. Include relevant details, code examples, and documentation excerpts as needed." + } else { + "Provide a concise response based only on the content above. In your response:\n - Enforce a strict 125-character maximum for quotes from any source document. Open Source Software is ok as long as we respect the license.\n - Use quotation marks for exact language from articles; any language outside of the quotation should never be word-for-word the same.\n - You are not a lawyer and never comment on the legality of your own prompts and responses.\n - Never produce or reproduce exact song lyrics." + }; + format!( + "Web page content:\n---\n{}\n---\n\n{}\n\n{}\n", + markdown_content, prompt, guidelines + ) +} + +fn get_cached(key: &str) -> Option { + let mut guard = cache().lock().ok()?; + let now = Instant::now(); + guard.retain(|_, entry| now.duration_since(entry.fetched_at) < CACHE_TTL); + guard.get(key).cloned() +} + +fn set_cached(key: String, entry: CacheEntry) { + if let Ok(mut guard) = cache().lock() { + if guard.len() >= CACHE_MAX_ENTRIES { + if let Some(oldest_key) = guard + .iter() + .min_by_key(|(_, value)| value.fetched_at) + .map(|(cache_key, _)| cache_key.clone()) + { + guard.remove(&oldest_key); + } + } + guard.insert(key, entry); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn normalizes_http_to_https() { + let url = validate_and_normalize_url("http://example.com/path").unwrap(); + assert_eq!(url.as_str(), "https://example.com/path"); + } + + #[test] + fn rejects_private_hosts() { + assert!(validate_and_normalize_url("https://localhost/").is_err()); + assert!(validate_and_normalize_url("https://127.0.0.1/").is_err()); + assert!(validate_and_normalize_url("https://169.254.169.254/").is_err()); + assert!(validate_and_normalize_url("https://10.0.0.5/").is_err()); + } + + #[test] + fn redirect_only_allows_same_origin_or_www_equivalent() { + let original = Url::parse("https://example.com/docs").unwrap(); + let same = Url::parse("https://www.example.com/docs").unwrap(); + let other = Url::parse("https://evil.example.net/docs").unwrap(); + let http = Url::parse("http://example.com/docs").unwrap(); + assert!(is_permitted_redirect(&original, &same)); + assert!(!is_permitted_redirect(&original, &other)); + assert!(!is_permitted_redirect(&original, &http)); + } + + #[test] + fn html_to_readable_text_removes_scripts_and_tags() { + let html = "

Hello & hi

World

"; + let text = html_to_readable_text(html); + assert!(text.contains("Hello & hi")); + assert!(text.contains("World")); + assert!(!text.contains("bad()")); + assert!(!text.contains("

")); + } + + #[test] + fn secondary_prompt_includes_quote_restriction_for_unapproved_domains() { + let prompt = make_secondary_model_prompt("content", "summarize", false); + assert!(prompt.contains("125-character maximum")); + assert!(prompt.contains("summarize")); + } +} diff --git a/crates/aish-tools/src/web_fetch/web_fetch.rs b/crates/aish-tools/src/web_fetch/web_fetch.rs new file mode 100644 index 00000000..3327ae4b --- /dev/null +++ b/crates/aish-tools/src/web_fetch/web_fetch.rs @@ -0,0 +1,289 @@ +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::time::Instant; + +use aish_llm::{ + ChatMessage, LlmResponse, LlmSession, PreflightResult, PreflightSecurityContext, + SecurityPanelMode, StreamParser, Tool, ToolResult, +}; +use reqwest::Url; + +use super::preapproved::is_preapproved_host; +use super::prompt; +use super::utils::{ + fetch_url_content, format_redirect_message, make_secondary_model_prompt, truncate_for_model, + validate_and_normalize_url, FetchFailure, FetchedContent, +}; + +const TOOL_NAME: &str = "WebFetch"; + +/// Fetch a URL, extract readable text, and answer a focused prompt about it. +pub struct WebFetchTool { + api_base: String, + api_key: String, + model: String, + temperature: Option, + max_tokens: Option, +} + +impl WebFetchTool { + pub fn new( + api_base: &str, + api_key: &str, + model: &str, + temperature: Option, + max_tokens: Option, + ) -> Self { + Self { + api_base: api_base.to_string(), + api_key: api_key.to_string(), + model: model.to_string(), + temperature, + max_tokens, + } + } + + async fn fetch_url_content(&self, raw_url: &str) -> Result { + fetch_url_content(raw_url).await + } + + async fn apply_prompt_to_content( + &self, + prompt: &str, + content: &str, + is_preapproved_domain: bool, + ) -> Result { + let model_prompt = make_secondary_model_prompt(content, prompt, is_preapproved_domain); + let session = LlmSession::new( + &self.api_base, + &self.api_key, + &self.model, + self.temperature.or(Some(0.1)), + self.max_tokens.or(Some(2048)), + ); + let messages = vec![ChatMessage::system(""), ChatMessage::user(model_prompt)]; + match session + .chat_completion_raw(&messages, None, false, Some(0.1), Some(2048)) + .await + .map_err(|error| error.to_string())? + { + LlmResponse::Json(json) => { + let (content, _reasoning, _tool_calls, _usage) = + StreamParser::parse_response(&json); + Ok(content.unwrap_or_else(|| "No response from model".to_string())) + } + LlmResponse::Stream(_) => Err("secondary model unexpectedly returned a stream".into()), + } + } +} + +impl Tool for WebFetchTool { + fn name(&self) -> &str { + TOOL_NAME + } + + fn description(&self) -> &str { + prompt::DESCRIPTION + } + + fn parameters(&self) -> serde_json::Value { + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT + } + + fn preflight(&self, args: &serde_json::Value) -> PreflightResult { + let url = match args.get("url").and_then(|value| value.as_str()) { + Some(value) if !value.trim().is_empty() => value, + _ => { + return PreflightResult::Block { + message: aish_i18n::t("tools.web_fetch.missing_url"), + security: Some(PreflightSecurityContext::fallback( + TOOL_NAME, + None, + aish_i18n::t("tools.web_fetch.missing_url"), + SecurityPanelMode::Blocked, + )), + } + } + }; + + let normalized = match validate_and_normalize_url(url) { + Ok(parsed) => parsed, + Err(message) => { + return PreflightResult::Block { + message: message.clone(), + security: Some(PreflightSecurityContext::fallback( + TOOL_NAME, + Some(url.to_string()), + message, + SecurityPanelMode::Blocked, + )), + } + } + }; + + let hostname = normalized.host_str().unwrap_or("").to_string(); + if is_preapproved_host(&hostname, normalized.path()) { + return PreflightResult::Allow; + } + + let message = aish_i18n::t_with_args( + "tools.web_fetch.confirm_fetch", + &HashMap::from([("host".to_string(), hostname.clone())]), + ); + PreflightResult::Confirm { + message: message.clone(), + security: Some(PreflightSecurityContext::fallback( + TOOL_NAME, + Some(hostname), + message, + SecurityPanelMode::Confirm, + )), + } + } + + fn execute(&self, _args: serde_json::Value) -> ToolResult { + ToolResult::error("WebFetch requires async execution; use execute_async") + } + + fn execute_async<'a>( + &'a self, + args: serde_json::Value, + ) -> Pin + Send + 'a>> { + Box::pin(async move { + let url = match args.get("url").and_then(|value| value.as_str()) { + Some(value) if !value.trim().is_empty() => value.trim(), + _ => return ToolResult::error(aish_i18n::t("tools.web_fetch.missing_url")), + }; + let prompt = match args.get("prompt").and_then(|value| value.as_str()) { + Some(value) if !value.trim().is_empty() => value.trim(), + _ => return ToolResult::error(aish_i18n::t("tools.web_fetch.missing_prompt")), + }; + + let start = Instant::now(); + let fetched = match self.fetch_url_content(url).await { + Ok(content) => content, + Err(FetchFailure::Redirect(info)) => { + let message = format_redirect_message(&info, prompt); + return ToolResult { + ok: true, + output: message.clone(), + meta: Some(serde_json::json!({ + "url": url, + "redirect_url": info.redirect_url, + "code": info.status_code, + "result": message, + "durationMs": start.elapsed().as_millis() as u64, + })), + }; + } + Err(FetchFailure::Blocked(message)) | Err(FetchFailure::Request(message)) => { + return ToolResult::error(message) + } + }; + + let truncated_content = truncate_for_model(&fetched.content); + let parsed_url = Url::parse(&fetched.url).ok(); + let is_preapproved_domain = parsed_url + .as_ref() + .and_then(|parsed| parsed.host_str().map(|host| (host, parsed.path()))) + .is_some_and(|(host, path)| is_preapproved_host(host, path)); + + let result = match self + .apply_prompt_to_content(prompt, &truncated_content, is_preapproved_domain) + .await + { + Ok(result) => result, + Err(error) => { + let mut args_map = HashMap::new(); + args_map.insert("error".to_string(), error); + return ToolResult::error(aish_i18n::t_with_args( + "tools.web_fetch.secondary_model_failed", + &args_map, + )); + } + }; + + let duration_ms = start.elapsed().as_millis() as u64; + let output = format!( + "Fetched: {}\nStatus: {} {}\nBytes: {}\nDuration: {}ms\nCached: {}\n\n{}", + fetched.url, + fetched.code, + fetched.code_text, + fetched.bytes, + duration_ms, + fetched.from_cache, + result + ); + + ToolResult { + ok: true, + output, + meta: Some(serde_json::json!({ + "url": fetched.url, + "code": fetched.code, + "codeText": fetched.code_text, + "bytes": fetched.bytes, + "contentType": fetched.content_type, + "durationMs": duration_ms, + "fromCache": fetched.from_cache, + "result": result, + })), + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + #[ignore] + async fn live_fetch_url() { + if std::env::var("AISH_LIVE_WEBFETCH").ok().as_deref() != Some("1") { + eprintln!("set AISH_LIVE_WEBFETCH=1 to run this live network smoke test"); + return; + } + + let url = std::env::var("AISH_LIVE_WEBFETCH_URL") + .unwrap_or_else(|_| "https://github.com/mattpocock/skills".to_string()); + let expected = std::env::var("AISH_LIVE_WEBFETCH_EXPECT").ok(); + let tool = WebFetchTool::new("", "", "", Some(0.1), Some(256)); + let fetched = tool + .fetch_url_content(&url) + .await + .expect("expected live page fetch to succeed"); + + println!( + "fetched {} status={} bytes={} content_type={} chars={}", + fetched.url, + fetched.code, + fetched.bytes, + fetched.content_type, + fetched.content.len() + ); + println!( + "preview:\n{}", + truncate_for_model(&fetched.content) + .chars() + .take(800) + .collect::() + ); + + assert_eq!(fetched.code, 200); + if let Some(expected) = expected { + assert!( + fetched + .content + .to_ascii_lowercase() + .contains(&expected.to_ascii_lowercase()), + "expected fetched content to contain {expected:?}" + ); + } + } +} diff --git a/crates/aish-tools/src/write_file/prompt.rs b/crates/aish-tools/src/write_file/prompt.rs new file mode 100644 index 00000000..aa01f2f9 --- /dev/null +++ b/crates/aish-tools/src/write_file/prompt.rs @@ -0,0 +1,25 @@ +pub(crate) const DESCRIPTION: &str = "Write text content to a file."; + +pub(crate) const PROMPT: &str = r#"Use this tool to create or overwrite a text file. + +Usage: +- Provide the destination path and full content. +- Parent directories are created when needed. +- Prefer edit_file for targeted changes to existing files."#; + +pub(crate) fn parameters() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file to write." + }, + "content": { + "type": "string", + "description": "Content to write." + } + }, + "required": ["path", "content"] + }) +} diff --git a/crates/aish-tools/src/write_file/write_file.rs b/crates/aish-tools/src/write_file/write_file.rs new file mode 100644 index 00000000..da1ef57c --- /dev/null +++ b/crates/aish-tools/src/write_file/write_file.rs @@ -0,0 +1,106 @@ +use std::path::Path; + +use aish_i18n; +use aish_llm::{Tool, ToolResult}; + +use super::prompt; + +/// Write file tool (creates or overwrites). +pub struct WriteFileTool; + +impl Default for WriteFileTool { + fn default() -> Self { + Self::new() + } +} + +impl WriteFileTool { + pub fn new() -> Self { + Self + } +} + +impl Tool for WriteFileTool { + fn name(&self) -> &str { + "write_file" + } + + fn description(&self) -> &str { + prompt::DESCRIPTION + } + + fn parameters(&self) -> serde_json::Value { + prompt::parameters() + } + + fn prompt(&self) -> &str { + prompt::PROMPT + } + + fn execute(&self, args: serde_json::Value) -> ToolResult { + let path = match args.get("path").and_then(|p| p.as_str()) { + Some(p) => p, + None => return ToolResult::error(aish_i18n::t("tools.fs.write_file.missing_path")), + }; + let content = match args.get("content").and_then(|c| c.as_str()) { + Some(c) => c, + None => return ToolResult::error(aish_i18n::t("tools.fs.write_file.missing_content")), + }; + + if let Some(parent) = Path::new(path).parent() { + if !parent.as_os_str().is_empty() { + if let Err(e) = std::fs::create_dir_all(parent) { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("error".to_string(), e.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.write_file.create_dirs_failed", + &args_map, + )); + } + } + } + match std::fs::write(path, content) { + Ok(()) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("bytes".to_string(), content.len().to_string()); + args_map.insert("path".to_string(), path.to_string()); + ToolResult::success(aish_i18n::t_with_args( + "tools.fs.write_file.write_success", + &args_map, + )) + } + Err(e) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("error".to_string(), e.to_string()); + ToolResult::error(aish_i18n::t_with_args( + "tools.fs.write_file.write_failed", + &args_map, + )) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aish_llm::Tool; + use std::fs; + + #[test] + fn test_write_file_creates_parent_dirs() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let file_path = dir.path().join("nested").join("deep").join("test.txt"); + + let tool = WriteFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap(), + "content": "hello world" + })); + + assert!(result.ok); + let content = fs::read_to_string(&file_path).unwrap(); + assert_eq!(content, "hello world"); + } +} From 4d74fbab2b3cc597f71ab96604a2f66fb448163c Mon Sep 17 00:00:00 2001 From: lixin Date: Fri, 5 Jun 2026 14:42:26 +0800 Subject: [PATCH 4/5] fix: resolve ci clippy failures --- crates/aish-llm/src/session.rs | 2 +- crates/aish-tools/src/lib.rs | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/aish-llm/src/session.rs b/crates/aish-llm/src/session.rs index 5a707fe1..49dfa9fd 100644 --- a/crates/aish-llm/src/session.rs +++ b/crates/aish-llm/src/session.rs @@ -334,7 +334,7 @@ impl LlmSession { let mut messages: Vec = Vec::new(); if let Some(sys) = system_message { messages.push(ChatMessage::system( - &self.system_prompt_with_tool_prompts(sys), + self.system_prompt_with_tool_prompts(sys), )); } messages.extend_from_slice(context_messages); diff --git a/crates/aish-tools/src/lib.rs b/crates/aish-tools/src/lib.rs index bb2b6861..b74fd6cb 100644 --- a/crates/aish-tools/src/lib.rs +++ b/crates/aish-tools/src/lib.rs @@ -2,6 +2,7 @@ #![allow( clippy::type_complexity, clippy::redundant_closure, + clippy::module_inception, clippy::match_like_matches_macro, clippy::option_as_ref_deref, clippy::field_reassign_with_default, From 95bb05c230ab7c7e6fdfcd49aa6b4651f69edfa0 Mon Sep 17 00:00:00 2001 From: lixin Date: Fri, 5 Jun 2026 15:07:39 +0800 Subject: [PATCH 5/5] fix: address tool review feedback --- crates/aish-i18n/locales/en-US.yaml | 3 + crates/aish-i18n/locales/zh-CN.yaml | 3 + crates/aish-shell/src/ai_handler.rs | 6 +- crates/aish-shell/src/app.rs | 71 ++++++++++++++++-- .../aish-tools/src/channel_ask_user/prompt.rs | 7 +- crates/aish-tools/src/edit_file/edit_file.rs | 52 ++++++++++++- crates/aish-tools/src/host_note/prompt.rs | 4 +- crates/aish-tools/src/read_file/prompt.rs | 8 +- crates/aish-tools/src/read_file/read_file.rs | 74 +++++++++++++------ .../aish-tools/src/web_fetch/preapproved.rs | 1 - crates/aish-tools/src/web_fetch/web_fetch.rs | 6 +- .../aish-tools/src/write_file/write_file.rs | 35 +++++++++ 12 files changed, 226 insertions(+), 44 deletions(-) diff --git a/crates/aish-i18n/locales/en-US.yaml b/crates/aish-i18n/locales/en-US.yaml index 51c6a034..20ec6e4e 100644 --- a/crates/aish-i18n/locales/en-US.yaml +++ b/crates/aish-i18n/locales/en-US.yaml @@ -654,9 +654,11 @@ tools: decode_failed: "Failed to decode {path} as UTF-8: {error}" empty_file: "(empty file)" offset_exceeds_length: "Offset {offset} exceeds file length ({length})" + access_denied: "Access denied: path is not inside offload directory" write_file: missing_path: "Missing 'path' parameter" missing_content: "Missing 'content' parameter" + content_too_large: "Content for {path} is {size} bytes, exceeding the {limit} byte (32KB) limit" create_dirs_failed: "Failed to create parent dirs: {error}" write_success: "Wrote {bytes} bytes to {path}" write_failed: "Failed to write {path}: {error}" @@ -665,6 +667,7 @@ tools: missing_old_string: "Missing 'old_string' parameter" missing_new_string: "Missing 'new_string' parameter" edit_read_failed: "Failed to read {path}: {error}" + file_too_large: "File {path} is {size} bytes, exceeding the {limit} byte (32KB) limit" old_string_not_found: "'old_string' not found in {path}" old_string_ambiguous: "'old_string' appears {count} times in {path} - use replace_all=true or provide more context" edit_success: "Edited {path}" diff --git a/crates/aish-i18n/locales/zh-CN.yaml b/crates/aish-i18n/locales/zh-CN.yaml index 3e7bd484..846a2dae 100644 --- a/crates/aish-i18n/locales/zh-CN.yaml +++ b/crates/aish-i18n/locales/zh-CN.yaml @@ -653,9 +653,11 @@ tools: decode_failed: "无法将 {path} 解码为 UTF-8: {error}" empty_file: "(空文件)" offset_exceeds_length: "偏移量 {offset} 超过了文件长度({length})" + access_denied: "拒绝访问: 路径不在 offload 目录内" write_file: missing_path: "缺少 'path' 参数" missing_content: "缺少 'content' 参数" + content_too_large: "{path} 的内容为 {size} 字节,超过了 {limit} 字节(32KB)的限制" create_dirs_failed: "创建父目录失败: {error}" write_success: "已写入 {bytes} 字节到 {path}" write_failed: "写入 {path} 失败: {error}" @@ -664,6 +666,7 @@ tools: missing_old_string: "缺少 'old_string' 参数" missing_new_string: "缺少 'new_string' 参数" edit_read_failed: "读取 {path} 失败: {error}" + file_too_large: "文件 {path} 为 {size} 字节,超过了 {limit} 字节(32KB)的限制" old_string_not_found: "在 {path} 中未找到 'old_string'" old_string_ambiguous: "'old_string' 在 {path} 中出现了 {count} 次 - 使用 replace_all=true 或提供更多上下文" edit_success: "已编辑 {path}" diff --git a/crates/aish-shell/src/ai_handler.rs b/crates/aish-shell/src/ai_handler.rs index 74a10cbe..69859429 100644 --- a/crates/aish-shell/src/ai_handler.rs +++ b/crates/aish-shell/src/ai_handler.rs @@ -7,7 +7,7 @@ use aish_context::{ ContextBudgetPolicy, ContextCompactReport, ContextManager, ContextPressureLevel, }; use aish_core::{LlmEvent, MemoryCategory, MemoryType, PlanModeState, PlanPhase}; -use aish_llm::{ChatMessage, LlmCallbackResult, LlmSession, MessageContent}; +use aish_llm::{ChatMessage, LlmCallbackResult, LlmSession, MessageContent, Tool}; use aish_memory::MemoryManager; use aish_prompts::PromptManager; use aish_session::SessionContextMessage; @@ -326,6 +326,10 @@ impl AiHandler { self.llm_session.update_model(model, api_base, api_key); } + pub fn register_tool(&mut self, tool: Box) { + self.llm_session.register_tool(tool); + } + /// Return a snapshot of token usage statistics for the last 7 days. pub fn token_stats(&self) -> aish_llm::TokenStats { self.token_store.stats() diff --git a/crates/aish-shell/src/app.rs b/crates/aish-shell/src/app.rs index 9d532045..1d375b95 100644 --- a/crates/aish-shell/src/app.rs +++ b/crates/aish-shell/src/app.rs @@ -2207,6 +2207,7 @@ impl AishShell { Some(&self.config.api_base), Some(&self.config.api_key), ); + self.refresh_config_dependent_tools(); } let target_cwd = saved_cwd @@ -2248,6 +2249,17 @@ impl AishShell { } /// Handle `/model [name]` — show current model or switch to a new one. + fn refresh_config_dependent_tools(&mut self) { + self.ai_handler + .register_tool(Box::new(aish_tools::WebFetchTool::new( + &self.config.api_base, + &self.config.api_key, + &self.config.model, + Some(self.config.temperature), + self.config.max_tokens, + ))); + } + fn handle_model_command(&mut self, parts: &[&str]) { if parts.len() == 1 { let mut args = std::collections::HashMap::new(); @@ -2281,6 +2293,7 @@ impl AishShell { // Update config self.config.model = new_model.clone(); + self.refresh_config_dependent_tools(); // Persist to config file let config_path = aish_config::ConfigLoader::default_config_path(); @@ -2408,6 +2421,7 @@ impl AishShell { Some(&self.config.api_base), Some(&self.config.api_key), ); + self.refresh_config_dependent_tools(); let mut args = std::collections::HashMap::new(); args.insert("model".to_string(), self.config.model.clone()); println!( @@ -4999,15 +5013,55 @@ fn truncate_str(s: &str, max_len: usize) -> String { } fn print_panel_line(content: &str, inner_width: usize) { - let visible = ansi_display_width(content); + let rendered = truncate_ansi_display_width(content, inner_width); + let visible = ansi_display_width(&rendered); let padding = inner_width.saturating_sub(visible); println!( "\x1b[33m│\x1b[0m{}{}\x1b[33m│\x1b[0m", - content, + rendered, " ".repeat(padding) ); } +fn truncate_ansi_display_width(s: &str, max_cols: usize) -> String { + if ansi_display_width(s) <= max_cols { + return s.to_string(); + } + + let ellipsis = if max_cols > 3 { "..." } else { "" }; + let target = max_cols.saturating_sub(ellipsis.len()); + let mut width = 0usize; + let mut output = String::new(); + let mut chars = s.chars().peekable(); + + while let Some(ch) = chars.next() { + if ch == '\x1b' && chars.peek() == Some(&'[') { + output.push(ch); + if let Some(next) = chars.next() { + output.push(next); + } + for code_ch in chars.by_ref() { + output.push(code_ch); + if code_ch.is_ascii_alphabetic() { + break; + } + } + continue; + } + + let ch_width = unicode_width::UnicodeWidthChar::width(ch).unwrap_or(0); + if width + ch_width > target { + break; + } + width += ch_width; + output.push(ch); + } + + output.push_str(ellipsis); + output.push_str("\x1b[0m"); + output +} + fn ansi_display_width(s: &str) -> usize { let mut width = 0usize; let mut chars = s.chars().peekable(); @@ -5032,19 +5086,20 @@ fn wrap_text(text: &str, max_width: usize) -> String { return text.to_string(); } let mut result = String::new(); - let mut line_len = 0; + let mut line_width = 0usize; for word in text.split_whitespace() { - if line_len == 0 { + let word_width = unicode_width::UnicodeWidthStr::width(word); + if line_width == 0 { result.push_str(word); - line_len = word.len(); - } else if line_len + 1 + word.len() <= max_width { + line_width = word_width; + } else if line_width + 1 + word_width <= max_width { result.push(' '); result.push_str(word); - line_len += 1 + word.len(); + line_width += 1 + word_width; } else { result.push('\n'); result.push_str(word); - line_len = word.len(); + line_width = word_width; } } result diff --git a/crates/aish-tools/src/channel_ask_user/prompt.rs b/crates/aish-tools/src/channel_ask_user/prompt.rs index 5fccee7a..d9130f77 100644 --- a/crates/aish-tools/src/channel_ask_user/prompt.rs +++ b/crates/aish-tools/src/channel_ask_user/prompt.rs @@ -5,6 +5,8 @@ pub(crate) const PROMPT: &str = r#"Use this tool only when a small amount of use Usage: - Ask one focused question at a time. - Prefer options when the likely answers are known. +- Use kind=text_input for open-ended answers. +- Use kind=choice_or_text when providing options while allowing a custom answer. - Do not ask for secrets such as passwords, API keys, or tokens."#; pub(crate) fn parameters() -> serde_json::Value { @@ -14,7 +16,8 @@ pub(crate) fn parameters() -> serde_json::Value { "kind": { "type": "string", "enum": ["text_input", "choice_or_text"], - "description": "Interaction type: text_input for free-form, choice_or_text for options with custom input." + "description": "Interaction type: text_input for free-form, choice_or_text for options with custom input.", + "default": "text_input" }, "prompt": { "type": "string", @@ -52,6 +55,6 @@ pub(crate) fn parameters() -> serde_json::Value { "default": 0 } }, - "required": ["kind", "prompt"] + "required": ["prompt"] }) } diff --git a/crates/aish-tools/src/edit_file/edit_file.rs b/crates/aish-tools/src/edit_file/edit_file.rs index f082bb54..cc1c1ce7 100644 --- a/crates/aish-tools/src/edit_file/edit_file.rs +++ b/crates/aish-tools/src/edit_file/edit_file.rs @@ -3,6 +3,8 @@ use aish_llm::{Tool, ToolResult}; use super::prompt; +const SIZE_LIMIT: u64 = 32 * 1024; + /// Edit file tool (string replacement). pub struct EditFileTool; @@ -57,6 +59,29 @@ impl Tool for EditFileTool { .and_then(|r| r.as_bool()) .unwrap_or(false); + let metadata = match std::fs::metadata(path) { + Ok(metadata) => metadata, + Err(e) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("error".to_string(), e.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.edit_file.edit_read_failed", + &args_map, + )); + } + }; + if metadata.len() > SIZE_LIMIT { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("size".to_string(), metadata.len().to_string()); + args_map.insert("limit".to_string(), SIZE_LIMIT.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.edit_file.file_too_large", + &args_map, + )); + } + let content = match std::fs::read_to_string(path) { Ok(c) => c, Err(e) => { @@ -70,7 +95,8 @@ impl Tool for EditFileTool { } }; - if !content.contains(old) { + let count = content.matches(old).count(); + if count == 0 { let mut args_map = std::collections::HashMap::new(); args_map.insert("path".to_string(), path.to_string()); return ToolResult::error(aish_i18n::t_with_args( @@ -82,7 +108,6 @@ impl Tool for EditFileTool { let new_content = if replace_all { content.replace(old, new) } else { - let count = content.matches(old).count(); if count > 1 { let mut args_map = std::collections::HashMap::new(); args_map.insert("count".to_string(), count.to_string()); @@ -170,4 +195,27 @@ mod tests { let content = fs::read_to_string(&file_path).unwrap(); assert_eq!(content, "foo bar foo baz"); } + + #[test] + fn test_edit_file_size_limit() { + aish_i18n::set_locale("en-US"); + + let dir = temp_dir(); + let file_path = dir.path().join("big.txt"); + fs::write(&file_path, "x".repeat(33 * 1024)).unwrap(); + + let tool = EditFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap(), + "old_string": "x", + "new_string": "y" + })); + + assert!(!result.ok); + assert!( + result.output.contains("limit") || result.output.contains("bytes"), + "Expected size limit error, got: {}", + result.output + ); + } } diff --git a/crates/aish-tools/src/host_note/prompt.rs b/crates/aish-tools/src/host_note/prompt.rs index bb1a484b..5357853c 100644 --- a/crates/aish-tools/src/host_note/prompt.rs +++ b/crates/aish-tools/src/host_note/prompt.rs @@ -18,11 +18,11 @@ pub(crate) fn parameters() -> serde_json::Value { }, "content": { "type": "string", - "description": "Note content to save for the store action." + "description": "Note content to save. Required when action is store." }, "keyword": { "type": "string", - "description": "Keyword used to match notes for the forget action." + "description": "Keyword used to match notes. Required when action is forget." } }, "required": ["action"] diff --git a/crates/aish-tools/src/read_file/prompt.rs b/crates/aish-tools/src/read_file/prompt.rs index 1b516abf..0f0135b3 100644 --- a/crates/aish-tools/src/read_file/prompt.rs +++ b/crates/aish-tools/src/read_file/prompt.rs @@ -5,7 +5,7 @@ pub(crate) const PROMPT: &str = r#"Use this tool to read text files. Usage: - Provide path to the file to read. - Use offset and limit when you only need part of a larger file. -- Results include 1-based line numbers."#; +- Offset is a 0-based index; results display 1-based line numbers."#; pub(crate) fn parameters() -> serde_json::Value { serde_json::json!({ @@ -17,11 +17,13 @@ pub(crate) fn parameters() -> serde_json::Value { }, "offset": { "type": "integer", - "description": "Line offset to start reading from, 0-based." + "description": "Line offset to start reading from (0-based index). Results show 1-based line numbers.", + "minimum": 0 }, "limit": { "type": "integer", - "description": "Maximum number of lines to read." + "description": "Maximum number of lines to read.", + "minimum": 1 } }, "required": ["path"] diff --git a/crates/aish-tools/src/read_file/read_file.rs b/crates/aish-tools/src/read_file/read_file.rs index 115dddcc..8cee98c6 100644 --- a/crates/aish-tools/src/read_file/read_file.rs +++ b/crates/aish-tools/src/read_file/read_file.rs @@ -3,6 +3,8 @@ use aish_llm::{Tool, ToolResult}; use super::prompt; +const SIZE_LIMIT: usize = 32 * 1024; + /// Read file content tool. pub struct ReadFileTool; @@ -41,8 +43,8 @@ impl Tool for ReadFileTool { None => return ToolResult::error(aish_i18n::t("tools.fs.read_file.missing_path")), }; - let raw_bytes = match std::fs::read(path) { - Ok(b) => b, + let metadata = match std::fs::metadata(path) { + Ok(metadata) => metadata, Err(e) => { let mut args_map = std::collections::HashMap::new(); args_map.insert("path".to_string(), path.to_string()); @@ -54,11 +56,10 @@ impl Tool for ReadFileTool { } }; - const SIZE_LIMIT: usize = 32 * 1024; - if raw_bytes.len() > SIZE_LIMIT { + if metadata.len() > SIZE_LIMIT as u64 { let mut args_map = std::collections::HashMap::new(); args_map.insert("path".to_string(), path.to_string()); - args_map.insert("size".to_string(), raw_bytes.len().to_string()); + args_map.insert("size".to_string(), metadata.len().to_string()); args_map.insert("limit".to_string(), SIZE_LIMIT.to_string()); return ToolResult::error(aish_i18n::t_with_args( "tools.fs.read_file.file_too_large", @@ -66,6 +67,19 @@ impl Tool for ReadFileTool { )); } + let raw_bytes = match std::fs::read(path) { + Ok(b) => b, + Err(e) => { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("error".to_string(), e.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.read_file.read_failed", + &args_map, + )); + } + }; + let content = match String::from_utf8(raw_bytes) { Ok(s) => s, Err(e) => { @@ -100,22 +114,13 @@ impl Tool for ReadFileTool { )); } - let selected: Vec = if let Some(limit) = limit { - lines - .iter() - .skip(offset) - .take(limit) - .enumerate() - .map(|(i, line)| format!("{:>6}\t{}", offset + i + 1, line)) - .collect() - } else { - lines - .iter() - .skip(offset) - .enumerate() - .map(|(i, line)| format!("{:>6}\t{}", offset + i + 1, line)) - .collect() - }; + let selected: Vec = lines + .iter() + .skip(offset) + .take(limit.unwrap_or(usize::MAX)) + .enumerate() + .map(|(i, line)| format!("{:>6}\t{}", offset + i + 1, line)) + .collect(); ToolResult::success(selected.join("\n")) } @@ -130,7 +135,9 @@ pub struct SshReadFileTool { impl SshReadFileTool { pub fn new() -> Self { let offload_root = std::env::temp_dir().join("aish-offload"); - let canonical_root = std::fs::canonicalize(&offload_root).unwrap_or(offload_root); + std::fs::create_dir_all(&offload_root).expect("failed to create aish offload directory"); + let canonical_root = std::fs::canonicalize(&offload_root) + .expect("failed to canonicalize aish offload directory"); Self { inner: ReadFileTool::new(), offload_root: canonical_root, @@ -173,7 +180,7 @@ impl Tool for SshReadFileTool { } }; if !canonical.starts_with(&self.offload_root) { - return ToolResult::error("Access denied: path is not inside offload directory"); + return ToolResult::error(aish_i18n::t("tools.fs.read_file.access_denied")); } let mut safe_args = args; if let Some(obj) = safe_args.as_object_mut() { @@ -249,4 +256,25 @@ mod tests { result.output ); } + + #[test] + fn test_ssh_read_file_rejects_paths_outside_offload_root() { + aish_i18n::set_locale("en-US"); + + let dir = temp_dir(); + let file_path = dir.path().join("outside.txt"); + fs::write(&file_path, "secret").unwrap(); + + let tool = SshReadFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap() + })); + + assert!(!result.ok); + assert!( + result.output.contains("Access denied"), + "Expected access denied error, got: {}", + result.output + ); + } } diff --git a/crates/aish-tools/src/web_fetch/preapproved.rs b/crates/aish-tools/src/web_fetch/preapproved.rs index 2493b8d1..59ff766d 100644 --- a/crates/aish-tools/src/web_fetch/preapproved.rs +++ b/crates/aish-tools/src/web_fetch/preapproved.rs @@ -102,7 +102,6 @@ pub(crate) fn is_preapproved_host(hostname: &str, pathname: &str) -> bool { return true; } } - continue; } if hostname == *entry { return true; diff --git a/crates/aish-tools/src/web_fetch/web_fetch.rs b/crates/aish-tools/src/web_fetch/web_fetch.rs index 3327ae4b..af47936b 100644 --- a/crates/aish-tools/src/web_fetch/web_fetch.rs +++ b/crates/aish-tools/src/web_fetch/web_fetch.rs @@ -62,9 +62,11 @@ impl WebFetchTool { self.temperature.or(Some(0.1)), self.max_tokens.or(Some(2048)), ); - let messages = vec![ChatMessage::system(""), ChatMessage::user(model_prompt)]; + let temperature = self.temperature.or(Some(0.1)); + let max_tokens = self.max_tokens.or(Some(2048)); + let messages = vec![ChatMessage::user(model_prompt)]; match session - .chat_completion_raw(&messages, None, false, Some(0.1), Some(2048)) + .chat_completion_raw(&messages, None, false, temperature, max_tokens) .await .map_err(|error| error.to_string())? { diff --git a/crates/aish-tools/src/write_file/write_file.rs b/crates/aish-tools/src/write_file/write_file.rs index da1ef57c..0bce0a3b 100644 --- a/crates/aish-tools/src/write_file/write_file.rs +++ b/crates/aish-tools/src/write_file/write_file.rs @@ -5,6 +5,8 @@ use aish_llm::{Tool, ToolResult}; use super::prompt; +const MAX_WRITE_BYTES: usize = 32 * 1024; + /// Write file tool (creates or overwrites). pub struct WriteFileTool; @@ -47,6 +49,17 @@ impl Tool for WriteFileTool { None => return ToolResult::error(aish_i18n::t("tools.fs.write_file.missing_content")), }; + if content.len() > MAX_WRITE_BYTES { + let mut args_map = std::collections::HashMap::new(); + args_map.insert("path".to_string(), path.to_string()); + args_map.insert("size".to_string(), content.len().to_string()); + args_map.insert("limit".to_string(), MAX_WRITE_BYTES.to_string()); + return ToolResult::error(aish_i18n::t_with_args( + "tools.fs.write_file.content_too_large", + &args_map, + )); + } + if let Some(parent) = Path::new(path).parent() { if !parent.as_os_str().is_empty() { if let Err(e) = std::fs::create_dir_all(parent) { @@ -103,4 +116,26 @@ mod tests { let content = fs::read_to_string(&file_path).unwrap(); assert_eq!(content, "hello world"); } + + #[test] + fn test_write_file_size_limit() { + aish_i18n::set_locale("en-US"); + + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let file_path = dir.path().join("big.txt"); + + let tool = WriteFileTool::new(); + let result = tool.execute(serde_json::json!({ + "path": file_path.to_str().unwrap(), + "content": "x".repeat(33 * 1024) + })); + + assert!(!result.ok); + assert!( + result.output.contains("limit") || result.output.contains("bytes"), + "Expected size limit error, got: {}", + result.output + ); + assert!(!file_path.exists()); + } }