新增: AI 工具(download_file + http output_file + grep 窗口 + fetch_search 搜索 + get_app_config + generate_image)
download_file 跨平台 URL 到文件流式下载;http_request output_file 落盘 + 截断可配; grep context_chars 大单行窗口;fetch_search DuckDuckGo 免 key 搜索;get_app_config 只读返当前配置(治 AI 查配置绕 PowerShell);generate_image SenseNova 图像生成。
This commit is contained in:
@@ -37,8 +37,11 @@ use super::super::secret::{
|
||||
// 提供商管理
|
||||
// ============================================================
|
||||
|
||||
/// api_key 脱敏:IPC 不传明文给前端(FR-S1),保留首尾各 4 字符便于辨识
|
||||
fn mask_api_key(key: &str) -> String {
|
||||
/// api_key 脱敏:IPC 不传明文给前端(FR-S1),保留首尾各 4 字符便于辨识。
|
||||
///
|
||||
/// `pub(crate)` 供 get_app_config AI 工具复用(查 DevFlow 自身 provider 配置时同样脱敏,
|
||||
/// 不向 LLM 回灌明文 key,与 ai_list_providers IPC 走同一套脱敏规则,单点维护)。
|
||||
pub(crate) fn mask_api_key(key: &str) -> String {
|
||||
let chars: Vec<char> = key.chars().collect();
|
||||
if chars.len() <= 8 {
|
||||
return "•".repeat(chars.len());
|
||||
|
||||
@@ -0,0 +1,308 @@
|
||||
//! download_file AI 工具 — URL → 文件(跨平台流式下载,无 shell/PowerShell 陷阱)
|
||||
//!
|
||||
//! 设计目标:让 LLM 把远程文件(URL)落到本地磁盘,替代 `run_command` 里的
|
||||
//! `curl -o` / `wget` / `Invoke-WebRequest`。后者在 Windows 上踩 alias、编码、
|
||||
//! 跨平台陷阱(PowerShell alias 表不同、GBK 解码、Invoke-WebRequest 不在 PS Core 外可用、
|
||||
//! curl 在 Win 是 alias = Invoke-WebRequest 而非真 curl)。本工具纯 reqwest + tokio::fs,
|
||||
//! 跨平台一致,无 shell spawn。
|
||||
//!
|
||||
//! ## 与 http_request / fetch_url 的分工
|
||||
//!
|
||||
//! - `http_request`:结构化 API 调用,返回 body(截断 50KB),不落盘。
|
||||
//! - `fetch_url`:网页 → markdown 文档嗅探,返回 markdown,不落盘。
|
||||
//! - `download_file`:**URL → 落盘文件**(二进制/大文件友好,流式写,不进 prompt)。
|
||||
//!
|
||||
//! ## 安全
|
||||
//!
|
||||
//! 1. **SSRF 防护**:复用 `commands::ai::http` 的 `validate_url` / `resolve_and_check_host` /
|
||||
//! `build_client` / `execute_with_redirects`(协议白名单 + 私网 IP 黑名单 + DNS resolve 后
|
||||
//! 校验防 rebinding + 重定向 ≤3 跳每跳重校验)。只发 GET,无 body。
|
||||
//! 2. **路径校验**:复用 `tool_registry::validate_path` + `resolve_workspace_path_with_allowed`
|
||||
//! (URL 解码防 %2e%2e 遍历 + 词法层越界兜底 + canonicalize symlink 逃逸拦截 + 授权目录白名单)。
|
||||
//! 与 write_file 同款,output_path 必须落在授权目录内。
|
||||
//! 3. **大小上限**:默认 200MB(可调 const),流式边下边累计,超限立即中止 + 删半成品,防磁盘炸。
|
||||
//! 4. **原子写**:tmp → rename(同目录,不跨卷),写一半崩溃不留半成品文件(对齐 write_file FR-S7)。
|
||||
//!
|
||||
//! ## 风险:High
|
||||
//!
|
||||
//! 下载文件有副作用(写磁盘 + 触发外发网络 + 可能拿到可执行/敏感内容),统一 High 须人工批准,
|
||||
//! 与 write_file(Medium)区别在于:download 的来源是外部不可信网络,不可预知内容。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::AsyncWriteExt;
|
||||
|
||||
use crate::commands::ai::http::{
|
||||
build_client, execute_with_redirects, resolve_and_check_host, validate_url, MAX_REDIRECTS,
|
||||
};
|
||||
use crate::commands::ai::tool_registry::{
|
||||
DEFAULT_RUN_COMMAND_TIMEOUT_SECS, MAX_RUN_COMMAND_TIMEOUT_SECS,
|
||||
};
|
||||
|
||||
/// 下载文件大小硬上限(字节)。200MB —— 防超大文件(ISO/镜像/视频)撑爆磁盘。
|
||||
/// 流式边下边累计,超限立即中止 + 删半成品。需更大文件走专门工具/分片,不在此。
|
||||
const MAX_DOWNLOAD_BYTES: u64 = 200 * 1024 * 1024;
|
||||
|
||||
/// download_file 工具 handler 入口(供 tools/download_file.rs register 调用)。
|
||||
///
|
||||
/// 参数:
|
||||
/// - url: 必填,http/https(SSRF 防护拒私网/localhost)
|
||||
/// - output_path: 必填,落盘路径(必须落在授权目录内,复用 write_file 同款路径校验)
|
||||
/// - timeout_secs: 可选,默认 DEFAULT_RUN_COMMAND_TIMEOUT_SECS(60),clamp ≤MAX(600)
|
||||
///
|
||||
/// 返回 {path, bytes_written, status, url}
|
||||
///
|
||||
/// 注:handler 不直接拿 allowed_dirs(tools/download_file.rs 闭包捕获并在调用前解析路径)。
|
||||
/// 本函数接收「已解析的绝对路径字符串」(resolved_path),路径校验在调用方完成(闭包持有 allowed_dirs,
|
||||
/// 解析后传字符串进来,与 write_file 模式一致 —— resolve_workspace_path_with_allowed 需 allowed_dirs
|
||||
/// 快照,故解析留在闭包内,本函数专注网络 + 落盘)。
|
||||
pub(crate) async fn execute_download_file(
|
||||
args: Value,
|
||||
resolved_output_path: String,
|
||||
) -> anyhow::Result<Value> {
|
||||
// ── 参数解析 ──
|
||||
let url_raw = args.get("url").and_then(|v| v.as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("缺少 url 参数"))?
|
||||
.trim()
|
||||
.to_string();
|
||||
if url_raw.is_empty() {
|
||||
anyhow::bail!("url 不能为空");
|
||||
}
|
||||
if resolved_output_path.trim().is_empty() {
|
||||
anyhow::bail!("output_path 不能为空");
|
||||
}
|
||||
let timeout_secs = args.get("timeout_secs").and_then(|v| v.as_u64())
|
||||
.unwrap_or(DEFAULT_RUN_COMMAND_TIMEOUT_SECS)
|
||||
.clamp(1, MAX_RUN_COMMAND_TIMEOUT_SECS);
|
||||
|
||||
// ── SSRF 校验(三层:词法 URL → DNS resolve IP → 每跳重定向重复) ──
|
||||
let (_scheme, host, port) = validate_url(&url_raw)?;
|
||||
resolve_and_check_host(&host, port).await?;
|
||||
|
||||
// ── 构建限制 client + 执行 GET(手动重定向循环复用 http.rs) ──
|
||||
let client = build_client(Duration::from_secs(timeout_secs))?;
|
||||
let started = Instant::now();
|
||||
let resp = execute_with_redirects(
|
||||
&client,
|
||||
reqwest::Method::GET,
|
||||
url_raw.clone(),
|
||||
&HashMap::new(), // download_file 不带自定义头(纯 GET 文件)
|
||||
&None, // 无 body
|
||||
MAX_REDIRECTS,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let status = resp.status().as_u16();
|
||||
let final_url = resp.url().to_string();
|
||||
if !resp.status().is_success() {
|
||||
anyhow::bail!(
|
||||
"下载失败:HTTP {} {}({})",
|
||||
status,
|
||||
resp.status().canonical_reason().unwrap_or(""),
|
||||
final_url
|
||||
);
|
||||
}
|
||||
|
||||
// ── 流式写文件(避免大文件 OOM) ──
|
||||
// response.bytes_stream() 需 reqwest "stream" feature(src-tauri/Cargo.toml 已声明)。
|
||||
// 边下边写 + 累计字节,超 MAX_DOWNLOAD_BYTES 立即中止 + 删半成品。
|
||||
let target = std::path::Path::new(&resolved_output_path);
|
||||
|
||||
// 父目录必须存在且在授权目录内(防 path=授权根时 parent 越界 create_dir_all)。
|
||||
// 路径白名单校验已在调用方(resolve_workspace_path_with_allowed)完成,这里只补 parent 授权 +
|
||||
// create_dir_all(对齐 write_file FR-S8)。parent 授权复检由调用方解析时已覆盖(resolved 已过白名单),
|
||||
// create_dir_all 仅针对尚不存在的 parent 目录创建。
|
||||
if let Some(parent) = target.parent() {
|
||||
tokio::fs::create_dir_all(parent).await
|
||||
.map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?;
|
||||
}
|
||||
|
||||
// 原子写:tmp → rename(同目录,不跨卷)。写一半崩溃/超限中止均删 tmp,不留半成品。
|
||||
let tmp = format!("{}.tmp-download", resolved_output_path);
|
||||
let mut file = tokio::fs::File::create(&tmp).await
|
||||
.map_err(|e| anyhow::anyhow!("创建临时文件失败 ({}): {}", tmp, e))?;
|
||||
|
||||
let mut stream = resp.bytes_stream();
|
||||
let mut written: u64 = 0;
|
||||
let mut exceeded = false;
|
||||
use futures::StreamExt;
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|e| anyhow::anyhow!("读取响应流失败: {}", e))?;
|
||||
written += chunk.len() as u64;
|
||||
if written > MAX_DOWNLOAD_BYTES {
|
||||
exceeded = true;
|
||||
break;
|
||||
}
|
||||
file.write_all(&chunk).await
|
||||
.map_err(|e| anyhow::anyhow!("写入文件失败: {}", e))?;
|
||||
}
|
||||
// flush + 关闭(file drop 前 flush,确保数据落盘再 rename)
|
||||
file.flush().await
|
||||
.map_err(|e| anyhow::anyhow!("flush 文件失败: {}", e))?;
|
||||
drop(file);
|
||||
|
||||
if exceeded {
|
||||
let _ = tokio::fs::remove_file(&tmp).await; // 删半成品
|
||||
anyhow::bail!(
|
||||
"下载超过 {} 字节上限,已中止并清理临时文件",
|
||||
MAX_DOWNLOAD_BYTES
|
||||
);
|
||||
}
|
||||
|
||||
// tmp → 目标(原子替换,同目录不跨卷)。覆盖既有文件时旧文件被替换。
|
||||
if let Err(e) = tokio::fs::rename(&tmp, &resolved_output_path).await {
|
||||
let _ = tokio::fs::remove_file(&tmp).await;
|
||||
return Err(anyhow::anyhow!("原子替换失败: {}", e));
|
||||
}
|
||||
|
||||
let elapsed_ms = started.elapsed().as_millis() as u64;
|
||||
|
||||
Ok(json!({
|
||||
"path": resolved_output_path,
|
||||
"bytes_written": written,
|
||||
"status": status,
|
||||
"url": final_url,
|
||||
"elapsed_ms": elapsed_ms,
|
||||
}))
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// 单元测试
|
||||
//
|
||||
// 覆盖:
|
||||
// ① SSRF 拦截私网(validate_url 拒 127.0.0.1/10.x/169.254):走 handler 入口早期拒绝,
|
||||
// 断言错误信息含关键词(对齐 http.rs / fetch_url.rs 测试策略)。
|
||||
// ② 路径越界拒:validate_path 纯函数,断言 .. / 敏感目录被拒。
|
||||
// 真实下载集成走 #[ignore](CI 无网时跳过,本地手跑)。
|
||||
// ============================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::commands::ai::tool_registry::validate_path;
|
||||
|
||||
// ── validate_path 纯函数:路径越界 / 敏感目录拒(可单测,零依赖) ──
|
||||
|
||||
#[test]
|
||||
fn test_validate_path_rejects_traversal() {
|
||||
// .. 段被拒(分段归一化,不误伤 my..file)
|
||||
assert!(validate_path("../etc/passwd").is_err());
|
||||
assert!(validate_path("foo/../../bar").is_err());
|
||||
assert!(validate_path("a/..").is_err());
|
||||
// URL 编码绕过被解码后拦截(%2e%2e = ..)
|
||||
assert!(validate_path("%2e%2e/secret").is_err());
|
||||
assert!(validate_path("%2e%2e%2fetc").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_path_accepts_normal_relative() {
|
||||
// 正常相对路径通过(my..file 不误伤)
|
||||
assert!(validate_path("downloads/file.zip").is_ok());
|
||||
assert!(validate_path("my..file.txt").is_ok());
|
||||
assert!(validate_path("a/b/c").is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_path_rejects_sensitive_system_dirs() {
|
||||
// .ssh / .aws / .gnupg 等敏感目录被拒(走 is_in_system_blacklist)
|
||||
assert!(validate_path("/root/.ssh/id_rsa").is_err());
|
||||
assert!(validate_path("C:\\Users\\x\\.aws\\credentials").is_err());
|
||||
}
|
||||
|
||||
// ── handler 参数边界 ──
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_missing_url_errors() {
|
||||
let args = json!({ "output_path": "/tmp/x" });
|
||||
let err = execute_download_file(args, "/tmp/x".to_string()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("缺少 url"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_empty_url_errors() {
|
||||
let args = json!({ "url": " ", "output_path": "/tmp/x" });
|
||||
let err = execute_download_file(args, "/tmp/x".to_string()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("url 不能为空"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_empty_output_path_errors() {
|
||||
let args = json!({ "url": "https://example.com/", "output_path": " " });
|
||||
let err = execute_download_file(args, " ".to_string()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("output_path 不能为空"));
|
||||
}
|
||||
|
||||
// ── SSRF 拦截:handler 入口 validate_url 拒私网/localhost/非 http 协议(不发请求) ──
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_rejects_localhost() {
|
||||
let args = json!({ "url": "http://localhost:8080/secret", "output_path": "/tmp/x" });
|
||||
let err = execute_download_file(args, "/tmp/x".to_string()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("localhost") || format!("{}", err).contains("SSRF"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_rejects_private_ip_loopback() {
|
||||
// 127.0.0.1 环回
|
||||
let args = json!({ "url": "http://127.0.0.1/admin", "output_path": "/tmp/x" });
|
||||
let err = execute_download_file(args, "/tmp/x".to_string()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("私网") || format!("{}", err).contains("SSRF"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_rejects_private_ip_rfc1918_10() {
|
||||
// 10.x RFC1918 A 类私网
|
||||
let args = json!({ "url": "http://10.0.0.1/", "output_path": "/tmp/x" });
|
||||
let err = execute_download_file(args, "/tmp/x".to_string()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("私网") || format!("{}", err).contains("SSRF"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_rejects_metadata_endpoint() {
|
||||
// 169.254.169.254 云元数据服务(SSRF 头号目标)
|
||||
let args = json!({ "url": "http://169.254.169.254/latest/meta-data/", "output_path": "/tmp/x" });
|
||||
let err = execute_download_file(args, "/tmp/x".to_string()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("私网") || format!("{}", err).contains("SSRF"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_rejects_non_http_scheme() {
|
||||
// file/ftp/data 等非 http 协议被拒
|
||||
let args = json!({ "url": "file:///etc/passwd", "output_path": "/tmp/x" });
|
||||
let err = execute_download_file(args, "/tmp/x".to_string()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("协议"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_timeout_clamped_no_panic() {
|
||||
// timeout_secs 超大值应 clamp 到 MAX,不 panic(用 localhost 触发 SSRF 早期拒绝,
|
||||
// 验证 clamp 在拒绝前不 panic)
|
||||
let args = json!({ "url": "http://localhost/", "output_path": "/tmp/x", "timeout_secs": 999999 });
|
||||
let _ = execute_download_file(args, "/tmp/x".to_string()).await; // 不 panic 即通过
|
||||
}
|
||||
|
||||
// ── 真实下载集成(#[ignore]:CI 无网时跳过,本地手跑) ──
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "需真实网络 + 临时目录,CI 无网时跳过:cargo test -- --ignored"]
|
||||
async fn integration_download_example_com() {
|
||||
// 下载 example.com 首页到临时文件,验证 happy path。
|
||||
let tmp_dir = std::env::temp_dir().join("devflow-download-test");
|
||||
tokio::fs::create_dir_all(&tmp_dir).await.ok();
|
||||
let out = tmp_dir.join("example.html");
|
||||
let out_str = out.to_string_lossy().to_string();
|
||||
let args = json!({ "url": "https://example.com/", "output_path": out_str });
|
||||
let result = execute_download_file(args, out_str.clone()).await.expect("下载应成功");
|
||||
assert_eq!(result["status"].as_u64().unwrap(), 200);
|
||||
assert!(result["bytes_written"].as_u64().unwrap() > 0);
|
||||
assert!(result["url"].as_str().unwrap().contains("example.com"));
|
||||
// 文件确实落盘
|
||||
let meta = tokio::fs::metadata(&out_str).await.expect("文件应存在");
|
||||
assert!(meta.len() > 0);
|
||||
// 清理
|
||||
let _ = tokio::fs::remove_file(&out_str).await;
|
||||
let _ = tokio::fs::remove_dir(&tmp_dir).await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,781 @@
|
||||
//! fetch_search AI 工具 — 搜索引擎查询(DuckDuckGo HTML 免 key,GET → 解析结果列表)
|
||||
//!
|
||||
//! 设计目标:让 LLM 拿到「搜索」能力。用户给"调研 X"时,AI 需先搜索找文档入口,
|
||||
//! 再 fetch_url 抓。当前无搜索工具,LLM 只能瞎猜 URL。本工具 GET DuckDuckGo HTML 端点
|
||||
//! (免 API key),解析结果列表(标题+URL+摘要),让 AI 据此挑入口用 fetch_url 深读。
|
||||
//!
|
||||
//! ## 与 fetch_url / http_request 的分工
|
||||
//!
|
||||
//! - `http_request`:结构化 API 调用(POST/鉴权/JSON),拿原始响应 body。
|
||||
//! - `fetch_url`:**已知 URL** → markdown 文档嗅探,LLM 拿到正文。
|
||||
//! - `fetch_search`:**未知 URL** → 搜索引擎查询,LLM 拿到候选 URL 列表 + 摘要,
|
||||
//! 再据结果挑入口用 fetch_url 深读。补「搜索」缺口。
|
||||
//!
|
||||
//! ## HTML 解析方式:手写字符串扫描(非 scraper/html5ever)
|
||||
//!
|
||||
//! DDG HTML 端点结构稳定且单一来源,引入 scraper(html5ever + selectors + cssparser,
|
||||
//! 重依赖 ~5MB 编译产物)不值。对齐 fetch_url::extract_title 的简单字符串扫描风格:
|
||||
//! 在 HTML 中找 `class="result__a"` 锚点 → 取 href + 文本;找 `result__snippet` → 取摘要。
|
||||
//! 解析纯函数化(extract_results),零网络依赖,单测用 fixture HTML 字符串验证。
|
||||
//!
|
||||
//! ## DDG 重定向 URL 解包
|
||||
//!
|
||||
//! DDG 把结果 URL 包成 `https://duckduckgo.com/l/?uddg=<percent-encoded real url>&rut=...`
|
||||
//! 重定向跳板(统计/反爬)。本工具解析 `uddg=` 查询参数,percent-decode 还原真实 URL,
|
||||
//! LLM 拿到的就是可直接 fetch_url 的目标地址。解析失败(无 uddg 参数 / 非 DDG 跳板 URL)
|
||||
//! 保留原 href 兜底(可能是 //duckduckgo.com/... 内部页,LLM 至少能看到)。
|
||||
//!
|
||||
//! ## 安全(SSRF 防护 — 与 http_request / fetch_url 共享同一套)
|
||||
//!
|
||||
//! 复用 `http.rs` 的 `validate_url` / `resolve_and_check_host` / `build_client` /
|
||||
//! `execute_with_redirects`:① 协议白名单(仅 http/https)② 私网 IP 黑名单 ③ DNS resolve
|
||||
//! 后校验 IP(防 rebinding)④ 重定向 ≤3 跳每跳重校验。fetch_search 只发 GET 到固定 DDG 端点。
|
||||
//!
|
||||
//! ## 反爬评估
|
||||
//!
|
||||
//! DDG HTML 端点(html.duckduckgo.com/html/)对带浏览器 UA 的请求通常友好(无 key/无 token)。
|
||||
//! 本工具默认带主流浏览器 UA(与 fetch_url 一致)。若反爬升级(403/空结果),不硬上 ——
|
||||
//! handler 把 HTTP 状态/空结果如实返回,LLM 据此判断(不改用代理/绕过,避免对抗反爬)。
|
||||
//! 间歇限流:单次工具调用是单 GET,无高频请求,不内置 rate limit(并发由 LLM loop 控制,
|
||||
//! 搜索是低频操作,不至触发频控)。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
/// DuckDuckGo HTML 搜索端点(免 API key,GET 返回结果列表 HTML)。
|
||||
const DDG_HTML_ENDPOINT: &str = "https://html.duckduckgo.com/html/";
|
||||
|
||||
/// 默认返回结果数。LLM 单次搜索 5 条足以挑入口(太多撑爆 context)。
|
||||
const DEFAULT_MAX_RESULTS: usize = 5;
|
||||
/// 返回结果数下限。防 LLM 传极小值拿到无意义单条。
|
||||
const MIN_MAX_RESULTS: usize = 1;
|
||||
/// 返回结果数上限。防 LLM 传极大值仍撑爆 context(10 条覆盖多数场景)。
|
||||
const MAX_MAX_RESULTS: usize = 10;
|
||||
|
||||
/// 默认请求超时(秒)。与 http_request / fetch_url 一致。
|
||||
const DEFAULT_TIMEOUT_SECS: u64 = 30;
|
||||
/// 超时硬上限(秒)。与 http_request / fetch_url 一致。
|
||||
const MAX_TIMEOUT_SECS: u64 = 60;
|
||||
|
||||
/// 响应 body 字节上限。DDG HTML 结果页通常 <500KB,2MB 足够且防 OOM(对齐 fetch_url)。
|
||||
const MAX_HTML_BYTES: usize = 2 * 1024 * 1024;
|
||||
|
||||
/// fetch_search 工具 handler 入口(供 tools/fetch_search.rs register 调用)。
|
||||
///
|
||||
/// 参数:
|
||||
/// - query: 必填,搜索关键词
|
||||
/// - max_results: 可选,返回结果数(默认 5,clamp [1, 10])
|
||||
///
|
||||
/// 返回 {query, results:[{title, url, snippet}], count}
|
||||
pub(crate) async fn execute_fetch_search(args: Value) -> anyhow::Result<Value> {
|
||||
// ── 参数解析 ──
|
||||
let query = args.get("query").and_then(|v| v.as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("缺少 query 参数"))?
|
||||
.trim()
|
||||
.to_string();
|
||||
if query.is_empty() {
|
||||
anyhow::bail!("query 不能为空");
|
||||
}
|
||||
let max_results = (args.get("max_results").and_then(|v| v.as_u64())
|
||||
.unwrap_or(DEFAULT_MAX_RESULTS as u64) as usize)
|
||||
.clamp(MIN_MAX_RESULTS, MAX_MAX_RESULTS);
|
||||
let timeout_secs = args.get("timeout_secs").and_then(|v| v.as_u64())
|
||||
.unwrap_or(DEFAULT_TIMEOUT_SECS)
|
||||
.min(MAX_TIMEOUT_SECS)
|
||||
.max(1);
|
||||
|
||||
// ── 构造请求 URL:percent-encode query 拼 DDG 端点 ?q=<encoded> ──
|
||||
// DDG HTML 端点接受 GET ?q=<url-encoded query>。encode 严格 RFC 3986(query 段)。
|
||||
let encoded_query = percent_encode_query(&query);
|
||||
let search_url = format!("{}?q={}", DDG_HTML_ENDPOINT, encoded_query);
|
||||
|
||||
// ── SSRF 校验(对固定 DDG 端点也走标准防护:协议/host/DNS resolve IP) ──
|
||||
let (_scheme, host, port) = crate::commands::ai::http::validate_url(&search_url)?;
|
||||
crate::commands::ai::http::resolve_and_check_host(&host, port).await?;
|
||||
|
||||
// ── 构建限制 client + 执行 GET(复用 http.rs 重定向循环,带浏览器 UA) ──
|
||||
// DDG HTML 端点对非浏览器 UA 可能返回简化页/拒,带主流 UA 拿完整结果列表(对齐 fetch_url)。
|
||||
let mut headers: HashMap<String, String> = HashMap::new();
|
||||
headers.insert(
|
||||
"User-Agent".to_string(),
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36".to_string(),
|
||||
);
|
||||
// Accept-Language:让 DDG 返回匹配用户语言偏好的结果(中英文兼顾)。
|
||||
headers.insert(
|
||||
"Accept-Language".to_string(),
|
||||
"en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7".to_string(),
|
||||
);
|
||||
|
||||
let client = crate::commands::ai::http::build_client(Duration::from_secs(timeout_secs))?;
|
||||
let started = Instant::now();
|
||||
let resp = crate::commands::ai::http::execute_with_redirects(
|
||||
&client,
|
||||
reqwest::Method::GET,
|
||||
search_url.clone(),
|
||||
&headers,
|
||||
&None,
|
||||
crate::commands::ai::http::MAX_REDIRECTS,
|
||||
)
|
||||
.await?;
|
||||
let elapsed_ms = started.elapsed().as_millis() as u64;
|
||||
|
||||
let final_url = resp.url().to_string();
|
||||
let status = resp.status().as_u16();
|
||||
if !resp.status().is_success() {
|
||||
anyhow::bail!(
|
||||
"搜索失败:HTTP {} {}({})。DDG 可能限流/反爬,稍后重试或换关键词",
|
||||
status,
|
||||
resp.status().canonical_reason().unwrap_or(""),
|
||||
final_url
|
||||
);
|
||||
}
|
||||
|
||||
// ── 读取 body + 大小挡板 ──
|
||||
let bytes = resp.bytes().await
|
||||
.map_err(|e| anyhow::anyhow!("读取响应 body 失败: {}", e))?;
|
||||
let total_bytes = bytes.len();
|
||||
if total_bytes > MAX_HTML_BYTES {
|
||||
anyhow::bail!(
|
||||
"响应过大:{} 字节超过 {} 上限,fetch_search 拿到异常大响应",
|
||||
total_bytes, MAX_HTML_BYTES
|
||||
);
|
||||
}
|
||||
let html = String::from_utf8_lossy(&bytes).into_owned();
|
||||
|
||||
// ── 解析结果列表(纯函数,单测用 fixture HTML 验证) ──
|
||||
let mut results = extract_results(&html);
|
||||
// 截断到 max_results(解析可能拿到更多,DDG 单页通常 ~30 条)
|
||||
results.truncate(max_results);
|
||||
|
||||
let count = results.len();
|
||||
tracing::info!(
|
||||
"fetch_search: query={:?} 拿到 {} 条结果(HTML {} 字节,{}ms)",
|
||||
query, count, total_bytes, elapsed_ms
|
||||
);
|
||||
|
||||
Ok(json!({
|
||||
"query": query,
|
||||
"results": results,
|
||||
"count": count,
|
||||
"status": status,
|
||||
"elapsed_ms": elapsed_ms,
|
||||
}))
|
||||
}
|
||||
|
||||
/// 从 DDG HTML 结果页解析结果列表(标题+URL+摘要)。
|
||||
///
|
||||
/// DDG HTML 每条结果结构(简化):
|
||||
/// ```text
|
||||
/// <div class="result ...">
|
||||
/// <h2 class="result__title">
|
||||
/// <a rel="nofollow" class="result__a" href="https://duckduckgo.com/l/?uddg=<encoded>&rut=...">Title</a>
|
||||
/// </h2>
|
||||
/// <a class="result__snippet" href="...">Snippet text...</a>
|
||||
/// </div>
|
||||
/// ```
|
||||
///
|
||||
/// 解析策略:扫描所有 `class="result__a"` 锚点 → 取 href + 内联文本;同序找
|
||||
/// `class="result__snippet"` → 取摘要。两者各扫一遍再按出现顺序对齐(DDG 标题与摘要
|
||||
/// 在每条结果内同序出现,故按位置一一对应;若数量不等则按较少者对齐)。
|
||||
///
|
||||
/// URL 解包:DDG 把真实 URL 包在 `uddg=<percent-encoded>` 查询参数(跳板重定向),
|
||||
/// 此处还原(percent-decode)让 LLM 拿到可直 fetch_url 的目标 URL。
|
||||
fn extract_results(html: &str) -> Vec<Value> {
|
||||
let titles_with_url: Vec<(String, String)> = extract_result_anchors(html);
|
||||
let snippets: Vec<String> = extract_snippets(html);
|
||||
|
||||
let n = titles_with_url.len().max(snippets.len());
|
||||
let mut results: Vec<Value> = Vec::with_capacity(n);
|
||||
for i in 0..n {
|
||||
let title_url = titles_with_url.get(i);
|
||||
let snippet = snippets.get(i);
|
||||
match (title_url, snippet) {
|
||||
(Some((title, href)), snip) => {
|
||||
let url = unwrap_ddg_redirect(href);
|
||||
results.push(json!({
|
||||
"title": title,
|
||||
"url": url,
|
||||
"snippet": snip.cloned().unwrap_or_default(),
|
||||
}));
|
||||
}
|
||||
(None, Some(snip)) => {
|
||||
// 残留 snippet 无对应标题(罕见,DDG 结构异常时兜底):跳过,无标题无 URL 无意义。
|
||||
// 不构造半残条目,保持每条都有 title+url(让 LLM 能 fetch_url)。
|
||||
}
|
||||
(None, None) => break,
|
||||
}
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
/// 提取所有 `class="result__a"` 锚点的 (内联文本, href),保持文档顺序。
|
||||
///
|
||||
/// 用字符串扫描(对齐 fetch_url::extract_title 风格),不引 scraper 重依赖。
|
||||
/// class 属性顺序不固定(DDG 偶有 `<a rel="nofollow" class="result__a" ...>` 或
|
||||
/// `<a class="result__a" ...>`),扫 `class="result__a"` 子串命中锚点起点,再从该点向后
|
||||
/// 找最近的 `<a ...>` 起标签尾(`>`),取 href 值 + 标签内文本(到 `</a>`)。
|
||||
fn extract_result_anchors(html: &str) -> Vec<(String, String)> {
|
||||
// MARKER 用 `result__a`(裸 class 名),命中后检查后续字符做边界判定:
|
||||
// 必须后跟非标识符字符(" 或空格,即 class 值结束或多个 class 间的分隔),避免 `result__a`
|
||||
// 子串误命中 `result__article` / `result__a_extra` 等同类 class(DDG 有 result__* 系列)。
|
||||
const MARKER: &str = "result__a";
|
||||
let marker_end_byte = MARKER.len();
|
||||
let mut out: Vec<(String, String)> = Vec::new();
|
||||
let mut search_from = 0usize;
|
||||
while search_from < html.len() {
|
||||
// 找下一个 "result__a" 出现位置(锚点 class 标记)
|
||||
let Some(rel_class) = html[search_from..].find(MARKER) else { break };
|
||||
let class_pos = search_from + rel_class;
|
||||
let after_marker = class_pos + marker_end_byte;
|
||||
// 边界检查:MARKER 后的字符必须是非标识符字符(" 或空格 或 >,即 class 值结束),
|
||||
// 否则可能是 `result__article` 这类以 `result__a` 为前缀的 class,跳过。
|
||||
let next_char = html[after_marker..].chars().next();
|
||||
let is_boundary = match next_char {
|
||||
Some(c) => !(c.is_ascii_alphanumeric() || c == '_' || c == '-'),
|
||||
None => true, // 字符串尾(罕见,视为边界)
|
||||
};
|
||||
if !is_boundary {
|
||||
// 非边界(如 result__article):跳过本命中继续找
|
||||
search_from = after_marker;
|
||||
continue;
|
||||
}
|
||||
|
||||
// 从 class 位置向前找最近的 `<a ` 起点(class 属性必在 <a ...> 标签内)
|
||||
// 倒扫到 class_pos 之前的最后一个 `<a `(`<a` 后接空白表示开标签)
|
||||
let head = &html[..class_pos];
|
||||
let Some(a_open_rel) = head.rfind("<a ") else {
|
||||
search_from = after_marker;
|
||||
continue;
|
||||
};
|
||||
// 从 a 开标签起点向后找标签尾 `>`
|
||||
let a_open_end_rel = html[a_open_rel..].find('>');
|
||||
let Some(a_open_end_offset) = a_open_end_rel else {
|
||||
search_from = after_marker;
|
||||
continue;
|
||||
};
|
||||
let a_open_end = a_open_rel + a_open_end_offset; // `>` 位置
|
||||
let a_open_tag = &html[a_open_rel..=a_open_end]; // 含 `<a ...>`
|
||||
|
||||
// 提取 href 值(href="..." 或 href='...')
|
||||
let href = extract_attr(a_open_tag, "href").unwrap_or_default();
|
||||
|
||||
// 取标签内文本(从 `>` 后到 `</a>`)
|
||||
let text_start = a_open_end + 1;
|
||||
let text_end = html[text_start..]
|
||||
.find("</a>")
|
||||
.map(|p| text_start + p)
|
||||
.unwrap_or(html.len());
|
||||
let raw_text = &html[text_start..text_end];
|
||||
let title = clean_html_text(raw_text);
|
||||
|
||||
if !title.is_empty() {
|
||||
out.push((title, href.to_string()));
|
||||
}
|
||||
|
||||
// 推进搜索位置到本锚点结束之后,避免重复命中同一锚点的 class
|
||||
search_from = text_end + "</a>".len();
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// 提取所有 `class="result__snippet"` 元素的内联文本,保持文档顺序。
|
||||
///
|
||||
/// DDG snippet 用 `<a class="result__snippet" ...>text</a>`(同时也是链接)。
|
||||
/// 扫 `result__snippet` 子串 → 向前找 `<a ` → 向后找 `>` → 取到 `</a>` 间文本。
|
||||
/// 与 extract_result_anchors 同款逻辑,仅 marker 不同。
|
||||
fn extract_snippets(html: &str) -> Vec<String> {
|
||||
const MARKER: &str = "result__snippet";
|
||||
let mut out: Vec<String> = Vec::new();
|
||||
let mut search_from = 0usize;
|
||||
while search_from < html.len() {
|
||||
let Some(rel) = html[search_from..].find(MARKER) else { break };
|
||||
let class_pos = search_from + rel;
|
||||
|
||||
let head = &html[..class_pos];
|
||||
let Some(a_open_rel) = head.rfind("<a ") else {
|
||||
search_from = class_pos + MARKER.len();
|
||||
continue;
|
||||
};
|
||||
let a_open_end_offset = match html[a_open_rel..].find('>') {
|
||||
Some(o) => o,
|
||||
None => {
|
||||
search_from = class_pos + MARKER.len();
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let a_open_end = a_open_rel + a_open_end_offset;
|
||||
|
||||
let text_start = a_open_end + 1;
|
||||
let text_end = html[text_start..]
|
||||
.find("</a>")
|
||||
.map(|p| text_start + p)
|
||||
.unwrap_or(html.len());
|
||||
let raw_text = &html[text_start..text_end];
|
||||
let snippet = clean_html_text(raw_text);
|
||||
if !snippet.is_empty() {
|
||||
out.push(snippet);
|
||||
}
|
||||
search_from = text_end + "</a>".len();
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// 从 `<a href="..." ...>` 等开标签中提取指定属性值。
|
||||
///
|
||||
/// 支持 `attr="..."` / `attr='...'`(单双引号),不支持无引号属性值(DDG href 总是带引号)。
|
||||
/// 值内的 HTML entity 不解码(简单场景 DDG href 不含 entity;若需可后续加 decode_html_entities)。
|
||||
///
|
||||
/// 已知局限:不做属性单词边界检查,故 `data-href="..."` 会被 `attr="href"` 匹配命中
|
||||
/// (因 "href=" 是 "data-href=" 的子串)。真实 DDG `<a>` 不同时带 data-href 与 href,
|
||||
/// 实务无影响;若未来需精确边界,改用前一个字符非字母数字的边界检查。
|
||||
fn extract_attr(open_tag: &str, attr: &str) -> Option<String> {
|
||||
let lower = open_tag.to_lowercase();
|
||||
let dq_pattern = format!("{}=\"", attr);
|
||||
let sq_pattern = format!("{}='", attr);
|
||||
// 优先双引号形式,其次单引号。matched_quote 记录命中模式对应的引号字符,
|
||||
// 避免 backtrack 推断(原 starts_with('"') 路径在 attr 出现在开标签首时 saturating_sub 边界不稳)。
|
||||
let (value_start_rel, quote_char): (usize, char) = lower
|
||||
.find(&dq_pattern)
|
||||
.map(|p| (p + dq_pattern.len(), '"'))
|
||||
.or_else(|| lower.find(&sq_pattern).map(|p| (p + sq_pattern.len(), '\'')))?;
|
||||
let value_end_rel = open_tag[value_start_rel..].find(quote_char)?;
|
||||
Some(open_tag[value_start_rel..value_start_rel + value_end_rel].to_string())
|
||||
}
|
||||
|
||||
/// 清理 HTML 内联文本:剥离嵌套标签(如 `<b>highlight</b>`)+ 折叠空白。
|
||||
///
|
||||
/// DDG 标题/snippet 内常有 `<b>` 高亮匹配关键词,需剥离得到纯文本。
|
||||
/// 不解码 entity(& 等),简单场景够用;后续如见 entity 可扩展。
|
||||
fn clean_html_text(raw: &str) -> String {
|
||||
// 剥离所有 `<...>` 子串(嵌套标签如 <b>word</b> → word)
|
||||
let mut stripped = String::with_capacity(raw.len());
|
||||
let mut in_tag = false;
|
||||
for c in raw.chars() {
|
||||
match c {
|
||||
'<' => in_tag = true,
|
||||
'>' => in_tag = false,
|
||||
_ if !in_tag => stripped.push(c),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
// 折叠连续空白(含 \n\t)为单空格,trim 首尾
|
||||
let cleaned: String = stripped.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
cleaned
|
||||
}
|
||||
|
||||
/// 解包 DDG 重定向跳板 URL 还原真实目标 URL。
|
||||
///
|
||||
/// DDG 把结果 URL 包成 `https://duckduckgo.com/l/?uddg=<percent-encoded real url>&rut=...`。
|
||||
/// 提取 `uddg` 查询参数值并 percent-decode → 真实 URL(可直接 fetch_url)。
|
||||
///
|
||||
/// 非 DDG 跳板 URL(无 `uddg=` 参数 / 解码失败)→ 原样返回 href 兜底。
|
||||
/// DDG 偶尔返回 `//duckduckgo.com/...` 协议相对内部页(无 uddg),保留原样让 LLM 至少看到。
|
||||
fn unwrap_ddg_redirect(href: &str) -> String {
|
||||
// 在 href 中找 `uddg=` 参数(可能在 ? 或 & 后)
|
||||
let Some(uddg_pos) = href.find("uddg=") else {
|
||||
return href.to_string();
|
||||
};
|
||||
let value_start = uddg_pos + "uddg=".len();
|
||||
let rest = &href[value_start..];
|
||||
// 取到下一个 `&` 或字符串尾
|
||||
let value_end = rest.find('&').unwrap_or(rest.len());
|
||||
let encoded = &rest[..value_end];
|
||||
// percent-decode(query 段,将 + 转空格 DDG 不用,uddg 值内真实 URL 的空格已 %20 编码)
|
||||
use percent_encoding::percent_decode_str;
|
||||
let decoded = percent_decode_str(encoded).decode_utf8_lossy().into_owned();
|
||||
if decoded.is_empty() {
|
||||
href.to_string()
|
||||
} else {
|
||||
decoded
|
||||
}
|
||||
}
|
||||
|
||||
/// percent-encode 搜索关键词(query 段 RFC 3986)。
|
||||
///
|
||||
/// 不引 urlencoding crate(http.rs 测试用同款手写),覆盖 query 段需转义字符:
|
||||
/// 空格 → %20、&=?# 等保留字符转义。字母/数字/常见安全符号(-._~)保留。
|
||||
fn percent_encode_query(s: &str) -> String {
|
||||
let mut out = String::with_capacity(s.len());
|
||||
for &b in s.as_bytes() {
|
||||
// RFC 3986 unreserved + 几个 query-safe 字符保留
|
||||
match b {
|
||||
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9'
|
||||
| b'-' | b'.' | b'_' | b'~' => out.push(b as char),
|
||||
_ => out.push_str(&format!("%{:02X}", b)),
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// 单元测试
|
||||
//
|
||||
// 覆盖纯函数(extract_results / extract_result_anchors / extract_snippets /
|
||||
// extract_attr / clean_html_text / unwrap_ddg_redirect / percent_encode_query),
|
||||
// 确定性零网络。handler 集成层走真实网络 #[ignore](对齐 http.rs / fetch_url 测试策略)。
|
||||
//
|
||||
// fixture HTML:基于 DDG HTML 端点真实结构简化(保留 result__a / result__snippet /
|
||||
// uddg 重定向跳板等关键标记),让解析逻辑可零网络验证。
|
||||
// ============================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// 真实 DDG HTML 结构简化 fixture(含 2 条结果,标题/snippet/uddg 重定向齐全)。
|
||||
/// 用于 extract_results 端到端解析验证(零网络)。
|
||||
const FIXTURE_DDG_HTML: &str = r#"<html><body>
|
||||
<div class="results">
|
||||
<div class="result results_links results_links_deep web-result">
|
||||
<h2 class="result__title">
|
||||
<a rel="nofollow" class="result__a" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Fdoc.rust-lang.org%2Fbook%2F&rut=abc">The Rust Programming Language <b>book</b></a>
|
||||
</h2>
|
||||
<a class="result__snippet" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Fdoc.rust-lang.org%2Fbook%2F">An introductory book about Rust covering syntax, ownership and <b>async</b>.</a>
|
||||
</div>
|
||||
<div class="result results_links results_links_deep web-result">
|
||||
<h2 class="result__title">
|
||||
<a rel="nofollow" class="result__a" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Ftokio.rs%2Ftutorial%2F&rut=def">Tokio Tutorial - Asynchronous Rust</a>
|
||||
</h2>
|
||||
<a class="result__snippet" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Ftokio.rs%2Ftutorial%2F">Learn how to write async Rust applications with the Tokio runtime.</a>
|
||||
</div>
|
||||
</div>
|
||||
</body></html>"#;
|
||||
|
||||
// ── extract_results:端到端解析(标题 + uddg 解包 URL + 摘要) ──
|
||||
|
||||
#[test]
|
||||
fn test_extract_results_end_to_end_two_results() {
|
||||
let results = extract_results(FIXTURE_DDG_HTML);
|
||||
assert_eq!(results.len(), 2, "应解析 2 条结果");
|
||||
|
||||
// 第 1 条
|
||||
let r0 = &results[0];
|
||||
assert_eq!(r0["title"].as_str().unwrap(), "The Rust Programming Language book");
|
||||
// uddg 解包后应是真实 URL,非 DDG 跳板
|
||||
assert_eq!(r0["url"].as_str().unwrap(), "https://doc.rust-lang.org/book/");
|
||||
assert!(r0["snippet"].as_str().unwrap().contains("introductory"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_results_second_result() {
|
||||
let results = extract_results(FIXTURE_DDG_HTML);
|
||||
let r1 = &results[1];
|
||||
assert_eq!(r1["title"].as_str().unwrap(), "Tokio Tutorial - Asynchronous Rust");
|
||||
assert_eq!(r1["url"].as_str().unwrap(), "https://tokio.rs/tutorial/");
|
||||
assert!(r1["snippet"].as_str().unwrap().contains("Tokio runtime"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_results_strips_b_tags_in_title() {
|
||||
// DDG 标题内常有 <b> 高亮匹配关键词,应剥离为纯文本
|
||||
let html = r#"<a class="result__a" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Fexample.com%2F">Hello <b>World</b></a>"#;
|
||||
let anchors = extract_result_anchors(html);
|
||||
assert_eq!(anchors.len(), 1);
|
||||
assert_eq!(anchors[0].0, "Hello World");
|
||||
}
|
||||
|
||||
// ── 截断 ──
|
||||
|
||||
#[test]
|
||||
fn test_execute_search_truncates_results() {
|
||||
// 不发网络:仅验证 extract_results + truncate 的截断逻辑(纯函数路径)
|
||||
// 构造 3 条 fixture,模拟 max_results=2 截断
|
||||
let html = r#"
|
||||
<a class="result__a" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Fa.com%2F">A</a>
|
||||
<a class="result__snippet" href="x">snip A</a>
|
||||
<a class="result__a" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Fb.com%2F">B</a>
|
||||
<a class="result__snippet" href="x">snip B</a>
|
||||
<a class="result__a" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Fc.com%2F">C</a>
|
||||
<a class="result__snippet" href="x">snip C</a>
|
||||
"#;
|
||||
let mut results = extract_results(html);
|
||||
results.truncate(2);
|
||||
assert_eq!(results.len(), 2, "应截断到 2 条");
|
||||
assert_eq!(results[0]["title"].as_str().unwrap(), "A");
|
||||
assert_eq!(results[1]["title"].as_str().unwrap(), "B");
|
||||
}
|
||||
|
||||
// ── extract_attr:开标签属性提取 ──
|
||||
|
||||
#[test]
|
||||
fn test_extract_attr_double_quote() {
|
||||
let tag = r#"<a href="https://example.com/" class="x">"#;
|
||||
assert_eq!(extract_attr(tag, "href").as_deref(), Some("https://example.com/"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_attr_single_quote() {
|
||||
let tag = r#"<a href='https://example.com/'>"#;
|
||||
assert_eq!(extract_attr(tag, "href").as_deref(), Some("https://example.com/"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_attr_missing_returns_none() {
|
||||
let tag = r#"<a class="x">"#;
|
||||
assert_eq!(extract_attr(tag, "href"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_attr_known_limitation_data_href() {
|
||||
// 已知局限记录:extract_attr 不做单词边界检查,"data-href=" 含 "href=" 子串会被命中。
|
||||
// 真实 DDG `<a>` 不同时带 data-href 与 href,实务无影响。本测试锁定当前行为供回归守护:
|
||||
// 若未来加边界检查改为不命中,此测试会红,提醒同步更新文档注释。
|
||||
let tag = r#"<a data-href="https://hidden.com/">"#;
|
||||
assert_eq!(
|
||||
extract_attr(tag, "href").as_deref(),
|
||||
Some("https://hidden.com/"),
|
||||
"已知局限:data-href 含 href= 子串会被命中(无单词边界检查)"
|
||||
);
|
||||
}
|
||||
|
||||
// ── clean_html_text:剥离嵌套标签 + 折叠空白 ──
|
||||
|
||||
#[test]
|
||||
fn test_clean_html_text_strips_tags() {
|
||||
assert_eq!(clean_html_text("Hello <b>World</b>!"), "Hello World!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clean_html_text_collapses_whitespace() {
|
||||
assert_eq!(clean_html_text(" multiple\n spaces\t here "), "multiple spaces here");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clean_html_text_empty_after_strip() {
|
||||
assert_eq!(clean_html_text("<b></b>"), "");
|
||||
assert_eq!(clean_html_text(" "), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clean_html_text_no_tags_passthrough() {
|
||||
assert_eq!(clean_html_text("plain text"), "plain text");
|
||||
}
|
||||
|
||||
// ── unwrap_ddg_redirect:uddg 参数解包 ──
|
||||
|
||||
#[test]
|
||||
fn test_unwrap_ddg_redirect_decodes_uddg_param() {
|
||||
let href = "https://duckduckgo.com/l/?uddg=https%3A%2F%2Fdoc.rust-lang.org%2F&rut=abc";
|
||||
assert_eq!(
|
||||
unwrap_ddg_redirect(href),
|
||||
"https://doc.rust-lang.org/"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unwrap_ddg_redirect_no_uddg_returns_original() {
|
||||
// 非 DDG 跳板(无 uddg 参数)→ 原样返回
|
||||
let href = "https://example.com/direct";
|
||||
assert_eq!(unwrap_ddg_redirect(href), "https://example.com/direct");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unwrap_ddg_redirect_uddg_last_param_no_amp() {
|
||||
// uddg 是最后一个参数(无尾随 &)→ 取到字符串尾
|
||||
let href = "https://duckduckgo.com/l/?uddg=https%3A%2F%2Ftokio.rs";
|
||||
assert_eq!(unwrap_ddg_redirect(href), "https://tokio.rs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unwrap_ddg_redirect_chinese_url_decoded() {
|
||||
// 中文 URL 经 percent-encode,decode 后应还原
|
||||
let href = "https://duckduckgo.com/l/?uddg=https%3A%2F%2Fzh.wikipedia.org%2Fwiki%2FRust&x=1";
|
||||
assert_eq!(
|
||||
unwrap_ddg_redirect(href),
|
||||
"https://zh.wikipedia.org/wiki/Rust"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unwrap_ddg_redirect_empty_encoded_returns_original() {
|
||||
// uddg= 空 → 解码空 → 返回原 href 兜底
|
||||
let href = "https://duckduckgo.com/l/?uddg=&rut=x";
|
||||
assert_eq!(unwrap_ddg_redirect(href), href);
|
||||
}
|
||||
|
||||
// ── percent_encode_query:搜索关键词编码 ──
|
||||
|
||||
#[test]
|
||||
fn test_percent_encode_query_space() {
|
||||
assert_eq!(percent_encode_query("rust async"), "rust%20async");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_percent_encode_query_chinese() {
|
||||
// 中文每字节转 %XX
|
||||
let encoded = percent_encode_query("Rust 语言");
|
||||
// "Rust " 保留 + "语"(3 字节 UTF-8)+ "言"(3 字节)
|
||||
assert!(encoded.starts_with("Rust%20"));
|
||||
assert!(encoded.contains("%E8%AF%AD")); // "语" 的 UTF-8 首三字节之一
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_percent_encode_query_safe_chars_preserved() {
|
||||
// unreserved 字母数字 -._~ 保留
|
||||
assert_eq!(percent_encode_query("abc-123_test.~"), "abc-123_test.~");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_percent_encode_query_special_chars_escaped() {
|
||||
// 保留字符 & = ? # 转义(防注入 query 解析)
|
||||
let encoded = percent_encode_query("a&b=c?d");
|
||||
assert_eq!(encoded, "a%26b%3Dc%3Fd");
|
||||
}
|
||||
|
||||
// ── extract_result_anchors:多结果顺序 + 跳过空标题 ──
|
||||
|
||||
#[test]
|
||||
fn test_extract_result_anchors_preserves_order() {
|
||||
let html = r#"
|
||||
<a class="result__a" href="https://x.com/1">First</a>
|
||||
<a class="result__a" href="https://x.com/2">Second</a>
|
||||
<a class="result__a" href="https://x.com/3">Third</a>
|
||||
"#;
|
||||
let anchors = extract_result_anchors(html);
|
||||
assert_eq!(anchors.len(), 3);
|
||||
assert_eq!(anchors[0].0, "First");
|
||||
assert_eq!(anchors[1].0, "Second");
|
||||
assert_eq!(anchors[2].0, "Third");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_result_anchors_skips_empty_title() {
|
||||
// 空标题(标签内仅空白)跳过,不构造半残条目
|
||||
let html = r#"
|
||||
<a class="result__a" href="https://x.com/1"> </a>
|
||||
<a class="result__a" href="https://x.com/2">Valid</a>
|
||||
"#;
|
||||
let anchors = extract_result_anchors(html);
|
||||
assert_eq!(anchors.len(), 1, "空标题应跳过");
|
||||
assert_eq!(anchors[0].0, "Valid");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_result_anchors_handles_rel_nofollow_first() {
|
||||
// DDG 真实顺序:`<a rel="nofollow" class="result__a" ...>`(rel 在 class 前)
|
||||
let html = r#"<a rel="nofollow" class="result__a" href="https://x.com/">Title</a>"#;
|
||||
let anchors = extract_result_anchors(html);
|
||||
assert_eq!(anchors.len(), 1);
|
||||
assert_eq!(anchors[0].1, "https://x.com/");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_result_anchors_no_results_returns_empty() {
|
||||
// 无 result__a 标记 → 空结果(无 panic)
|
||||
let html = "<html><body>no results here</body></html>";
|
||||
let anchors = extract_result_anchors(html);
|
||||
assert!(anchors.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_result_anchors_boundary_check_rejects_result_article() {
|
||||
// 边界检查:`result__article`(以 result__a 为前缀)不应被误命中为 result__a。
|
||||
// 防止 DDG 其他 result__* class(如 result__article)污染标题列表。
|
||||
let html = r#"
|
||||
<a class="result__article" href="https://wrong.com/">Should Not Match</a>
|
||||
<a class="result__a" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Fright.com%2F">Right Title</a>
|
||||
"#;
|
||||
let anchors = extract_result_anchors(html);
|
||||
assert_eq!(anchors.len(), 1, "result__article 不应被误命中");
|
||||
assert_eq!(anchors[0].0, "Right Title");
|
||||
assert_eq!(anchors[0].1, "https://duckduckgo.com/l/?uddg=https%3A%2F%2Fright.com%2F");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_result_anchors_handles_multi_class_value() {
|
||||
// class 列表含多个 class(`class="result__a result__url"`)仍应命中(result__a 后是空格)
|
||||
let html = r#"<a class="result__a extra" href="https://x.com/">Multi Class Title</a>"#;
|
||||
let anchors = extract_result_anchors(html);
|
||||
assert_eq!(anchors.len(), 1);
|
||||
assert_eq!(anchors[0].0, "Multi Class Title");
|
||||
}
|
||||
|
||||
// ── extract_snippets ──
|
||||
|
||||
#[test]
|
||||
fn test_extract_snippets_basic() {
|
||||
let html = r#"
|
||||
<a class="result__snippet" href="x">First snippet</a>
|
||||
<a class="result__snippet" href="y">Second snippet</a>
|
||||
"#;
|
||||
let snippets = extract_snippets(html);
|
||||
assert_eq!(snippets.len(), 2);
|
||||
assert_eq!(snippets[0], "First snippet");
|
||||
assert_eq!(snippets[1], "Second snippet");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_snippets_strips_b_tags() {
|
||||
let html = r#"<a class="result__snippet" href="x">Text with <b>highlight</b> word</a>"#;
|
||||
let snippets = extract_snippets(html);
|
||||
assert_eq!(snippets.len(), 1);
|
||||
assert_eq!(snippets[0], "Text with highlight word");
|
||||
}
|
||||
|
||||
// ── extract_results:无 snippet 时空字符串 ──
|
||||
|
||||
#[test]
|
||||
fn test_extract_results_snippet_missing_defaults_empty() {
|
||||
// 标题有但 snippet 缺(解析时 snippet.get(i) 返 None → 空字符串)
|
||||
let html = r#"
|
||||
<a class="result__a" href="https://duckduckgo.com/l/?uddg=https%3A%2F%2Fx.com%2F">Title Only</a>
|
||||
"#;
|
||||
let results = extract_results(html);
|
||||
assert_eq!(results.len(), 1);
|
||||
assert_eq!(results[0]["title"].as_str().unwrap(), "Title Only");
|
||||
assert_eq!(results[0]["snippet"].as_str().unwrap(), "", "无 snippet 应默认空字符串");
|
||||
}
|
||||
|
||||
// ── handler 参数边界(不发网络,触发早期拒绝) ──
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_missing_query_errors() {
|
||||
let args = json!({ "max_results": 5 });
|
||||
let err = execute_fetch_search(args).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("缺少 query"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_empty_query_errors() {
|
||||
let args = json!({ "query": " " });
|
||||
let err = execute_fetch_search(args).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("query 不能为空"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_max_results_clamped_low() {
|
||||
// max_results=0 应 clamp 到 1(用空 query 触发早期拒绝,验证 clamp 不 panic)
|
||||
let args = json!({ "query": " ", "max_results": 0 });
|
||||
let _ = execute_fetch_search(args).await; // 不 panic 即通过
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_max_results_clamped_high() {
|
||||
let args = json!({ "query": " ", "max_results": 999 });
|
||||
let _ = execute_fetch_search(args).await; // 不 panic 即通过
|
||||
}
|
||||
|
||||
// ── 真实网络集成(#[ignore]:CI 无网时跳过,本地手跑) ──
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "需真实网络(DDG),CI 无网时跳过:cargo test -- --ignored"]
|
||||
async fn integration_search_rust_async_returns_results() {
|
||||
let args = json!({ "query": "rust async programming", "max_results": 3 });
|
||||
let result = execute_fetch_search(args).await.expect("搜索应成功");
|
||||
let count = result["count"].as_u64().unwrap_or(0);
|
||||
assert!(count >= 1, "至少应返回 1 条结果,实际 {}", count);
|
||||
let results = result["results"].as_array().expect("results 应为数组");
|
||||
// 每条应有非空 title + url
|
||||
for r in results {
|
||||
assert!(!r["title"].as_str().unwrap_or("").is_empty(), "title 不应为空");
|
||||
let url = r["url"].as_str().unwrap_or("");
|
||||
assert!(!url.is_empty(), "url 不应为空");
|
||||
// uddg 应已解包,非 DDG 跳板 URL(除非 DDG 返回内部页)
|
||||
// 不强制断言非 duckduckgo.com(偶尔 DDG 返回内部结果),但应能 fetch_url
|
||||
}
|
||||
// query 应回显
|
||||
assert_eq!(result["query"].as_str().unwrap(), "rust async programming");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,587 @@
|
||||
//! generate_image AI 工具 — 调 provider 的 OpenAI 兼容图像生成端点生成图片并下载落地。
|
||||
//!
|
||||
//! 设计目标:让 LLM 调用图像生成模型(默认 SenseNova U1 Fast,也支持 DALL-E 等 OpenAI 兼容
|
||||
//! 图像端点)生成图片并落到工作区磁盘。图像生成模型不是 chat 模型,走 `/v1/images/generations`
|
||||
//! 端点(非 chat completions),故不能经主对话链路调用,必须用本工具显式触发。
|
||||
//!
|
||||
//! ## 与 http_request / fetch_url / download_file 的分工
|
||||
//!
|
||||
//! - `http_request`:结构化 API 调用,返回 body(截断 50KB),不落盘 —— 通用 HTTP,不感知图像端点。
|
||||
//! - `fetch_url`:网页 → markdown 文档嗅探,不落盘。
|
||||
//! - `download_file`:URL → 落盘文件(已知 URL)。
|
||||
//! - `generate_image`:**调 provider 图像端点 + 自动下载落地**(封装图像生成的 provider 凭证
|
||||
//! 解析 + 端点拼接 + 响应解析 + 图片下载全链路,LLM 只需给 prompt/model/output_path)。
|
||||
//!
|
||||
//! ## 风险:High
|
||||
//!
|
||||
//! ① 付费 API(图像生成按张计费);② 写磁盘(落盘 output_path);③ 外发网络(provider 端点
|
||||
//! + 图片 URL 下载)。统一 High 须人工批准。
|
||||
//!
|
||||
//! ## 安全
|
||||
//!
|
||||
//! 1. **provider 凭证**:从 db.ai_providers 筛 `enabled && provider_type=="openai_compat"` 的
|
||||
//! provider,经 `df_storage::secret::resolve_provider_secret` 解析 key(DB 优先 fallback keyring)。
|
||||
//! anthropic provider 不参与(图像端点是 OpenAI 风格)。
|
||||
//! 2. **端点拼接**:`build_images_url` 智能 base_url(`/v1`/`/v4` 后缀直接补 `/images/generations`,
|
||||
//! 否则补 `/v1/images/generations`),对齐 `model_fetch_helpers::build_models_url` 思路。
|
||||
//! 3. **provider 端点域名**:provider 已知域名(api.sensenova / api.openai / ...),非用户输入,
|
||||
//! SSRF 风险低,故 provider POST 直接 reqwest::Client 不走 SSRF 防护。
|
||||
//! 4. **图片 URL 下载**:图片 URL 来自 provider 响应,**可能被恶意 provider 篡改指向内网**
|
||||
//! (provider 域名虽可信但响应内容不可信),故下载图片 URL 复用 SSRF 防护
|
||||
//! (validate_url + resolve_and_check_host + build_client,与 download_file 同源)。
|
||||
//! 5. **路径校验**:output_path 走 `resolve_workspace_path_with_allowed`(白名单 + symlink 逃逸
|
||||
//! 拦截 + 路径遍历拒),与 write_file/download_file 同源;默认路径锚定首个持久授权目录。
|
||||
//! 6. **大小上限**:50MB(图片够用),流式边下边累计,超限中止 + 删半成品(对齐 download_file)。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use base64::{engine::general_purpose::STANDARD, Engine as _};
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use df_storage::crud::AiProviderRepo;
|
||||
use df_storage::db::Database;
|
||||
use df_storage::models::AiProviderRecord;
|
||||
use df_storage::secret::resolve_provider_secret;
|
||||
|
||||
use crate::commands::ai::http::{
|
||||
build_client, execute_with_redirects, resolve_and_check_host, validate_url, MAX_REDIRECTS,
|
||||
};
|
||||
use crate::commands::ai::tool_registry::resolve_workspace_path_with_allowed;
|
||||
use crate::state::AllowedDirs;
|
||||
|
||||
/// 下载图片大小硬上限(字节)。50MB —— 覆盖高分辨率图片(4K壁纸级 < 20MB),超限中止。
|
||||
/// 图片 URL 来自 provider 响应可能被篡改指向超大文件,流式边下边累计防磁盘炸。
|
||||
const MAX_IMAGE_BYTES: u64 = 50 * 1024 * 1024;
|
||||
|
||||
/// 生成图片 POST 请求超时(秒)。图像生成模型耗时较高(高分辨率 10-30s),给 120s 余量。
|
||||
const GENERATE_TIMEOUT_SECS: u64 = 120;
|
||||
|
||||
/// generate_image 工具 handler 入口(供 tools/generate_image.rs register 调用)。
|
||||
///
|
||||
/// 参数:
|
||||
/// - prompt: 必填,图片描述(中英文均可,空串报错)
|
||||
/// - model: 默认 `sensenova-u1-fast`(SenseNova U1 信息图生成,走 /v1/images/generations)
|
||||
/// - size: 可选,图片尺寸字符串(SenseNova U1 常用 `2752x1536`/`1536x2752`,DALL-E 常用 `1024x1024`)
|
||||
/// - n: 默认 1,生成数量,clamp [1, 4](多张图取首张 url/b64 落地,余下仅返回 url 列表)
|
||||
/// - provider_id: 可选,指定 provider(多 provider 时);不传则按 model 名匹配厂商 host
|
||||
/// - output_path: 可选,落盘路径;不传则默认 `{首个持久授权目录}/generated_images/{model}-{时间戳}.png`
|
||||
///
|
||||
/// 返回 {path, url(或 null 若仅 b64), model, provider_id, bytes_written, elapsed_ms}
|
||||
///
|
||||
/// 注:handler 接收 db + allowed_dirs(经 declare_tool! 闭包捕获 tuple Arc 传入)。
|
||||
/// 路径校验在本函数内完成(需 allowed_dirs 快照,与 http_request output_file 同款)。
|
||||
pub(crate) async fn execute_generate_image(
|
||||
args: Value,
|
||||
db: &Arc<Database>,
|
||||
allowed_dirs: &Arc<RwLock<AllowedDirs>>,
|
||||
) -> anyhow::Result<Value> {
|
||||
let started = Instant::now();
|
||||
|
||||
// ── 参数解析 ──
|
||||
let prompt = args.get("prompt").and_then(|v| v.as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("缺少 prompt 参数"))?
|
||||
.trim()
|
||||
.to_string();
|
||||
if prompt.is_empty() {
|
||||
anyhow::bail!("prompt 不能为空");
|
||||
}
|
||||
let model = args.get("model").and_then(|v| v.as_str())
|
||||
.unwrap_or("sensenova-u1-fast")
|
||||
.trim()
|
||||
.to_string();
|
||||
if model.is_empty() {
|
||||
anyhow::bail!("model 不能为空");
|
||||
}
|
||||
let size = args.get("size").and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty());
|
||||
let n = args.get("n").and_then(|v| v.as_u64())
|
||||
.unwrap_or(1)
|
||||
.clamp(1, 4) as u32;
|
||||
let provider_id = args.get("provider_id").and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty());
|
||||
let output_path_raw = args.get("output_path").and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty());
|
||||
|
||||
// ── 选 provider(enabled && openai_compat,按 model 名匹配厂商 host 兜底首个) ──
|
||||
let providers = AiProviderRepo::new(db).list_all().await?;
|
||||
let provider = select_image_provider(&providers, &model, provider_id.as_deref())?;
|
||||
let provider_id_resolved = provider.id.clone();
|
||||
|
||||
// ── 解析 api_key(DB 优先 fallback keyring,空则报错) ──
|
||||
let api_key = resolve_provider_secret(&provider);
|
||||
let api_key = api_key.trim();
|
||||
if api_key.is_empty() {
|
||||
anyhow::bail!(
|
||||
"provider「{}」未配置 api_key,请在设置填入并保存",
|
||||
provider.name
|
||||
);
|
||||
}
|
||||
|
||||
// ── 拼端点 URL ──
|
||||
let endpoint = build_images_url(&provider.base_url);
|
||||
|
||||
// ── POST 请求(provider 域名非用户输入,SSRF 风险低,直接 reqwest) ──
|
||||
let body = {
|
||||
let mut m = serde_json::Map::new();
|
||||
m.insert("model".into(), json!(model));
|
||||
m.insert("prompt".into(), json!(prompt));
|
||||
m.insert("n".into(), json!(n));
|
||||
if let Some(sz) = &size {
|
||||
m.insert("size".into(), json!(sz));
|
||||
}
|
||||
Value::Object(m)
|
||||
};
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(GENERATE_TIMEOUT_SECS))
|
||||
.connect_timeout(Duration::from_secs(15))
|
||||
.build()
|
||||
.map_err(|e| anyhow::anyhow!("HTTP client 构建失败: {}", e))?;
|
||||
let resp = client.post(&endpoint)
|
||||
.header("Authorization", format!("Bearer {}", api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("调用图像生成端点失败 ({}): {}", endpoint, e))?;
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
// 错误响应读 body 摘要(截断 500 chars)帮助定位
|
||||
let err_text = resp.text().await.unwrap_or_default();
|
||||
let snippet: String = err_text.chars().take(500).collect();
|
||||
anyhow::bail!(
|
||||
"图像生成端点返回 HTTP {}: {}",
|
||||
status.as_u16(),
|
||||
snippet
|
||||
);
|
||||
}
|
||||
let resp_body: Value = resp.json().await
|
||||
.map_err(|e| anyhow::anyhow!("解析图像生成响应 JSON 失败: {}", e))?;
|
||||
|
||||
// ── 解析响应:优先 data[0].url,次选 data[0].b64_json ──
|
||||
let data0 = resp_body.get("data").and_then(|d| d.get(0))
|
||||
.ok_or_else(|| {
|
||||
let snippet = body_snippet(&resp_body);
|
||||
anyhow::anyhow!("图像生成响应缺 data[0](响应摘要: {})", snippet)
|
||||
})?;
|
||||
let image_url_opt = data0.get("url").and_then(|u| u.as_str())
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.map(|s| s.trim().to_string());
|
||||
let b64_opt = data0.get("b64_json").and_then(|b| b.as_str())
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.map(|s| s.trim().to_string());
|
||||
|
||||
if image_url_opt.is_none() && b64_opt.is_none() {
|
||||
let snippet = body_snippet(&resp_body);
|
||||
anyhow::bail!(
|
||||
"图像生成响应 data[0] 既无 url 也无 b64_json(响应摘要: {})",
|
||||
snippet
|
||||
);
|
||||
}
|
||||
|
||||
// ── 解析 output_path(LLM 传则走白名单校验;不传则默认 generated_images/ 时间戳名) ──
|
||||
let snap = allowed_dirs.read().await.clone();
|
||||
let final_path = if let Some(raw) = &output_path_raw {
|
||||
resolve_workspace_path_with_allowed(raw, &snap)?
|
||||
.to_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("output_path 含非法字符"))?
|
||||
.to_string()
|
||||
} else {
|
||||
// 默认:首个持久授权目录/generated_images/{model}-{YYYYMMDD-HHMMSS}.png
|
||||
let root = snap.first_persistent_dir()
|
||||
.ok_or_else(|| anyhow::anyhow!(
|
||||
"未传 output_path 且无持久授权目录,请先绑定项目目录或显式传 output_path"
|
||||
))?;
|
||||
let ts = chrono::Local::now().format("%Y%m%d-%H%M%S");
|
||||
// model 含 / 等路径分隔符风险:替换为 - 防穿透到子目录
|
||||
let safe_model = model.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|'], "-");
|
||||
let fname = format!("{}-{}.png", safe_model, ts);
|
||||
root.join("generated_images").join(fname)
|
||||
.to_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("默认 output_path 含非法字符"))?
|
||||
.to_string()
|
||||
};
|
||||
|
||||
// ── 落盘:url 走 SSRF 防护下载流式写;b64 直接写解码后字节 ──
|
||||
let written_bytes = if let Some(image_url) = &image_url_opt {
|
||||
download_image_with_ssrf(image_url, &final_path).await?
|
||||
} else {
|
||||
// b64 路径:解码后直接写
|
||||
let b64 = b64_opt.as_ref().unwrap();
|
||||
let decoded = STANDARD.decode(b64)
|
||||
.map_err(|e| anyhow::anyhow!("b64_json 解码失败: {}", e))?;
|
||||
if (decoded.len() as u64) > MAX_IMAGE_BYTES {
|
||||
anyhow::bail!(
|
||||
"b64_json 解码后 {} 字节超过 {} 上限",
|
||||
decoded.len(),
|
||||
MAX_IMAGE_BYTES
|
||||
);
|
||||
}
|
||||
write_atomic(&final_path, &decoded).await?;
|
||||
decoded.len() as u64
|
||||
};
|
||||
|
||||
let elapsed_ms = started.elapsed().as_millis() as u64;
|
||||
|
||||
Ok(json!({
|
||||
"path": final_path,
|
||||
"url": image_url_opt,
|
||||
"model": model,
|
||||
"provider_id": provider_id_resolved,
|
||||
"bytes_written": written_bytes,
|
||||
"elapsed_ms": elapsed_ms,
|
||||
}))
|
||||
}
|
||||
|
||||
/// 从 provider 列表筛选图像生成可用 provider。
|
||||
///
|
||||
/// 筛选条件:`enabled && provider_type == "openai_compat"`(图像端点是 OpenAI 风格,
|
||||
/// anthropic provider 排除)。
|
||||
///
|
||||
/// 选择策略:
|
||||
/// 1. LLM 传 `provider_id` → 按 id 精确匹配(找不到报错)。
|
||||
/// 2. 否则按 model 名匹配厂商 host 提示:
|
||||
/// - model 含 `sensenova` → 找 base_url 含 `sensenova`
|
||||
/// - model 含 `dall-e` → 找 base_url 含 `openai`
|
||||
/// - model 含 `google`/`imagen` → 找 base_url 含 `googleapis`
|
||||
/// 3. host 命中 → 用首个匹配;host 未命中 → 兜底取筛选结果首个。
|
||||
/// 4. 筛选后空 → 报错(无可用 openai_compat provider)。
|
||||
fn select_image_provider(
|
||||
providers: &[AiProviderRecord],
|
||||
model: &str,
|
||||
provider_id: Option<&str>,
|
||||
) -> anyhow::Result<AiProviderRecord> {
|
||||
// 筛 enabled + openai_compat
|
||||
let candidates: Vec<&AiProviderRecord> = providers.iter()
|
||||
.filter(|p| p.enabled && p.provider_type == "openai_compat")
|
||||
.collect();
|
||||
|
||||
// 路径1:LLM 指定 provider_id,精确匹配
|
||||
if let Some(pid) = provider_id {
|
||||
let found = candidates.iter().find(|p| p.id == pid)
|
||||
.ok_or_else(|| anyhow::anyhow!(
|
||||
"未找到指定 provider(id={},且需 enabled + openai_compat),请检查 provider_id 参数",
|
||||
pid
|
||||
))?;
|
||||
return Ok((*found).clone());
|
||||
}
|
||||
|
||||
if candidates.is_empty() {
|
||||
anyhow::bail!("无可用 openai_compat provider,请先在设置配置支持图像生成的 provider");
|
||||
}
|
||||
|
||||
// 路径2:按 model 名匹配厂商 host
|
||||
let model_lower = model.to_lowercase();
|
||||
let host_hint: Option<&str> = if model_lower.contains("sensenova") {
|
||||
Some("sensenova")
|
||||
} else if model_lower.contains("dall-e") {
|
||||
Some("openai")
|
||||
} else if model_lower.contains("google") || model_lower.contains("imagen") {
|
||||
Some("googleapis")
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if let Some(hint) = host_hint {
|
||||
let hit = candidates.iter()
|
||||
.find(|p| p.base_url.to_lowercase().contains(hint));
|
||||
if let Some(p) = hit {
|
||||
return Ok((*p).clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 路径3:兜底首个候选
|
||||
Ok(candidates[0].clone())
|
||||
}
|
||||
|
||||
/// 拼接图像生成端点 URL。
|
||||
///
|
||||
/// 规则(base_url 先去尾 `/`,对齐 `model_fetch_helpers::build_models_url` 思路,
|
||||
/// 但目标是 `/images/generations`):
|
||||
/// - 以 `/v1` 或 `/v4` 结尾 → 直接补 `/images/generations`(避免 `/v1/v1/images/...`)
|
||||
/// - 否则 → 补 `/v1/images/generations`(最常见:用户填 `https://api.sensenova.cn`)
|
||||
fn build_images_url(base_url: &str) -> String {
|
||||
let url = base_url.trim_end_matches('/');
|
||||
if url.ends_with("/v1") || url.ends_with("/v4") {
|
||||
format!("{url}/images/generations")
|
||||
} else {
|
||||
format!("{url}/v1/images/generations")
|
||||
}
|
||||
}
|
||||
|
||||
/// 流式下载图片 URL 到 output_path,带 SSRF 防护 + 大小上限 + 原子写。
|
||||
///
|
||||
/// 图片 URL 来自 provider 响应可能被恶意 provider 篡改指向内网,故复用 SSRF 防护
|
||||
/// (validate_url + resolve_and_check_host + build_client + execute_with_redirects)。
|
||||
async fn download_image_with_ssrf(image_url: &str, output_path: &str) -> anyhow::Result<u64> {
|
||||
// SSRF 三层校验:词法 URL → DNS resolve IP → 每跳重定向重复
|
||||
let (_scheme, host, port) = validate_url(image_url)?;
|
||||
resolve_and_check_host(&host, port).await?;
|
||||
|
||||
let client = build_client(Duration::from_secs(GENERATE_TIMEOUT_SECS))?;
|
||||
let resp = execute_with_redirects(
|
||||
&client,
|
||||
reqwest::Method::GET,
|
||||
image_url.to_string(),
|
||||
&HashMap::new(), // 图片下载不带自定义头
|
||||
&None,
|
||||
MAX_REDIRECTS,
|
||||
)
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
anyhow::bail!(
|
||||
"下载图片失败:HTTP {} {}({})",
|
||||
resp.status().as_u16(),
|
||||
resp.status().canonical_reason().unwrap_or(""),
|
||||
image_url
|
||||
);
|
||||
}
|
||||
|
||||
// 流式写文件(避免大图 OOM),边下边累计,超 MAX_IMAGE_BYTES 立即中止 + 删半成品。
|
||||
// 父目录不存在自动创建。原子写:tmp → rename(同目录不跨卷)。
|
||||
let target = std::path::Path::new(output_path);
|
||||
if let Some(parent) = target.parent() {
|
||||
tokio::fs::create_dir_all(parent).await
|
||||
.map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?;
|
||||
}
|
||||
|
||||
let tmp = format!("{}.tmp-img", output_path);
|
||||
let mut file = tokio::fs::File::create(&tmp).await
|
||||
.map_err(|e| anyhow::anyhow!("创建临时文件失败 ({}): {}", tmp, e))?;
|
||||
|
||||
let mut stream = resp.bytes_stream();
|
||||
let mut written: u64 = 0;
|
||||
let mut exceeded = false;
|
||||
use futures::StreamExt;
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|e| anyhow::anyhow!("读取图片流失败: {}", e))?;
|
||||
written += chunk.len() as u64;
|
||||
if written > MAX_IMAGE_BYTES {
|
||||
exceeded = true;
|
||||
break;
|
||||
}
|
||||
file.write_all(&chunk).await
|
||||
.map_err(|e| anyhow::anyhow!("写入图片失败: {}", e))?;
|
||||
}
|
||||
file.flush().await
|
||||
.map_err(|e| anyhow::anyhow!("flush 图片文件失败: {}", e))?;
|
||||
drop(file);
|
||||
|
||||
if exceeded {
|
||||
let _ = tokio::fs::remove_file(&tmp).await;
|
||||
anyhow::bail!(
|
||||
"图片超过 {} 字节上限,已中止并清理临时文件",
|
||||
MAX_IMAGE_BYTES
|
||||
);
|
||||
}
|
||||
|
||||
if let Err(e) = tokio::fs::rename(&tmp, output_path).await {
|
||||
let _ = tokio::fs::remove_file(&tmp).await;
|
||||
return Err(anyhow::anyhow!("原子替换失败: {}", e));
|
||||
}
|
||||
|
||||
Ok(written)
|
||||
}
|
||||
|
||||
/// 原子写字节到指定路径(tmp → rename,同目录不跨卷)。
|
||||
/// b64 路径用此函数(已全部解码在内存,无需流式)。
|
||||
async fn write_atomic(output_path: &str, data: &[u8]) -> anyhow::Result<()> {
|
||||
let target = std::path::Path::new(output_path);
|
||||
if let Some(parent) = target.parent() {
|
||||
tokio::fs::create_dir_all(parent).await
|
||||
.map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?;
|
||||
}
|
||||
let tmp = format!("{}.tmp-img", output_path);
|
||||
tokio::fs::write(&tmp, data).await
|
||||
.map_err(|e| {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
anyhow::anyhow!("写入图片失败: {}", e)
|
||||
})?;
|
||||
tokio::fs::rename(&tmp, output_path).await
|
||||
.map_err(|e| {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
anyhow::anyhow!("原子替换失败: {}", e)
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 响应 body 截断为最多 500 chars 的字符串(用于错误信息),供定位用。
|
||||
fn body_snippet(v: &Value) -> String {
|
||||
let s = v.to_string();
|
||||
s.chars().take(500).collect()
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// 单元测试
|
||||
//
|
||||
// 覆盖:
|
||||
// ① build_images_url:端点拼接规则(纯函数,零依赖)
|
||||
// ② select_image_provider:provider 筛选逻辑(纯函数,内存构造 record)
|
||||
// ③ handler 参数边界(prompt 缺失/空、provider_id 不存在)
|
||||
// 真实 API 调用走 #[ignore](CI 无凭证/无网时跳过,本地手跑)。
|
||||
// ============================================================
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── build_images_url:端点拼接规则 ──
|
||||
|
||||
#[test]
|
||||
fn build_images_url_v1_suffix() {
|
||||
// /v1 结尾 → 直接补 /images/generations
|
||||
assert_eq!(
|
||||
build_images_url("https://api.openai.com/v1"),
|
||||
"https://api.openai.com/v1/images/generations"
|
||||
);
|
||||
// 不应产生 /v1/v1/images/...
|
||||
assert!(!build_images_url("https://api.openai.com/v1").contains("/v1/v1/"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_images_url_v4_suffix() {
|
||||
// GLM v4 风格 base
|
||||
assert_eq!(
|
||||
build_images_url("https://open.bigmodel.cn/api/paas/v4"),
|
||||
"https://open.bigmodel.cn/api/paas/v4/images/generations"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_images_url_bare_base() {
|
||||
// 用户只填根 → 补 /v1/images/generations
|
||||
assert_eq!(
|
||||
build_images_url("https://api.sensenova.cn"),
|
||||
"https://api.sensenova.cn/v1/images/generations"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_images_url_strips_trailing_slash() {
|
||||
assert_eq!(
|
||||
build_images_url("https://api.x.com/v1/"),
|
||||
"https://api.x.com/v1/images/generations"
|
||||
);
|
||||
}
|
||||
|
||||
// ── select_image_provider:筛选 + 匹配逻辑 ──
|
||||
|
||||
/// 构造测试用 AiProviderRecord(只填筛选关键字段)。
|
||||
fn mk_provider(id: &str, enabled: bool, ptype: &str, base_url: &str) -> AiProviderRecord {
|
||||
AiProviderRecord {
|
||||
id: id.into(),
|
||||
name: id.into(),
|
||||
provider_type: ptype.into(),
|
||||
api_key: String::new(),
|
||||
base_url: base_url.into(),
|
||||
default_model: String::new(),
|
||||
models: None,
|
||||
model_configs: Vec::new(),
|
||||
is_default: false,
|
||||
config: None,
|
||||
created_at: "0".into(),
|
||||
updated_at: "0".into(),
|
||||
enabled,
|
||||
weight: 50,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_provider_no_candidates_errors() {
|
||||
// 全 disabled 或非 openai_compat → 报错
|
||||
let providers = vec![
|
||||
mk_provider("p1", false, "openai_compat", "https://x"),
|
||||
mk_provider("p2", true, "anthropic", "https://y"),
|
||||
];
|
||||
let err = select_image_provider(&providers, "any-model", None).unwrap_err();
|
||||
assert!(format!("{}", err).contains("无可用"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_provider_provider_id_match() {
|
||||
let providers = vec![
|
||||
mk_provider("p1", true, "openai_compat", "https://api.sensenova.cn"),
|
||||
mk_provider("p2", true, "openai_compat", "https://api.openai.com"),
|
||||
];
|
||||
let got = select_image_provider(&providers, "any", Some("p2")).unwrap();
|
||||
assert_eq!(got.id, "p2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_provider_provider_id_not_found_errors() {
|
||||
let providers = vec![mk_provider("p1", true, "openai_compat", "https://x")];
|
||||
let err = select_image_provider(&providers, "any", Some("nonexistent")).unwrap_err();
|
||||
assert!(format!("{}", err).contains("未找到指定 provider"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_provider_host_match_sensenova() {
|
||||
// model 含 sensenova → 匹配 base_url 含 sensenova 的 provider
|
||||
let providers = vec![
|
||||
mk_provider("openai", true, "openai_compat", "https://api.openai.com"),
|
||||
mk_provider("sense", true, "openai_compat", "https://api.sensenova.cn"),
|
||||
];
|
||||
let got = select_image_provider(&providers, "sensenova-u1-fast", None).unwrap();
|
||||
assert_eq!(got.id, "sense");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_provider_host_match_dalle() {
|
||||
// model 含 dall-e → 匹配 base_url 含 openai 的 provider
|
||||
let providers = vec![
|
||||
mk_provider("glm", true, "openai_compat", "https://open.bigmodel.cn"),
|
||||
mk_provider("oai", true, "openai_compat", "https://api.openai.com"),
|
||||
];
|
||||
let got = select_image_provider(&providers, "dall-e-3", None).unwrap();
|
||||
assert_eq!(got.id, "oai");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_provider_host_unmatched_falls_back_first() {
|
||||
// model 名无厂商提示 → 兜底取筛选结果首个
|
||||
let providers = vec![
|
||||
mk_provider("first", true, "openai_compat", "https://api.unknown.com"),
|
||||
mk_provider("second", true, "openai_compat", "https://api.other.com"),
|
||||
];
|
||||
let got = select_image_provider(&providers, "my-custom-model", None).unwrap();
|
||||
assert_eq!(got.id, "first");
|
||||
}
|
||||
|
||||
// ── handler 参数边界(走 db 需 in-memory db) ──
|
||||
|
||||
#[tokio::test]
|
||||
async fn handler_missing_prompt_errors() {
|
||||
// 缺 prompt → 参数解析阶段报错(不发请求)
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
|
||||
let args = json!({ "model": "sensenova-u1-fast" });
|
||||
let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("prompt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handler_empty_prompt_errors() {
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
|
||||
let args = json!({ "prompt": " " });
|
||||
let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("prompt 不能为空"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handler_no_provider_errors() {
|
||||
// 无 openai_compat provider → 报错(不发请求)
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
|
||||
let args = json!({ "prompt": "一只猫" });
|
||||
let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("无可用") || format!("{}", err).contains("provider"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,383 @@
|
||||
//! get_app_config AI 工具 — 查 DevFlow 自身当前生效的 AI 配置(只读 + 脱敏)。
|
||||
//!
|
||||
//! ## 治什么症状
|
||||
//!
|
||||
//! 实测 AI 想看「当前 provider/model/agent 设置」时,因没原生「查配置」工具,只能 run_command
|
||||
//! 执行 PowerShell 内联 node/python 脚本查 db(node sqlite/...),命令行引号三重嵌套必失败
|
||||
//! (PowerShell 单引号包 node -e 内含 JS 模板字符串/双引号/db 路径,转义层数爆炸,复盘实证)。
|
||||
//! 本工具直接返当前配置 JSON,LLM 一调即得,根治 run_command 绕行。
|
||||
//!
|
||||
//! ## 安全(只读 + 脱敏)
|
||||
//!
|
||||
//! - **只读**:纯 list_all/load 内存值,无任何写库/写文件副作用。RiskLevel::Low。
|
||||
//! - **api_key 脱敏**:复用 [`mask_api_key`](crate::commands::ai::commands::mask_api_key)
|
||||
//! (前 4 + `••••` + 后 4),迁移后 DB api_key 空列 → 经 keyring 解析真实 key 再脱敏,
|
||||
//! 绝不向 LLM 回灌明文 key(对齐 ai_list_providers IPC 脱敏口径,单点维护)。
|
||||
//! - **app_settings 不返敏感 KV**:仅返与 AI 行为直接相关的已知无害 key
|
||||
//! (knowledge/timeout/custom_prompt 存在性 + 长度),自定义指令原文不回灌防指令注入。
|
||||
//!
|
||||
//! ## 数据来源(全部经 GetAppConfigCtx 捕获句柄,无 AppState 自引用)
|
||||
//!
|
||||
//! - providers / default_provider:`AiProviderRepo::new(&ctx.db).list_all()` + keyring 解析
|
||||
//! - agent:`ctx.agent_max_iterations` 等内存原子量当前 load 值(热改即反映)
|
||||
//! - app_settings:`SettingsRepo::new(&ctx.db).get_all()` 后白名单筛选
|
||||
//!
|
||||
//! ## 为何不直接捕获 &AppState
|
||||
//!
|
||||
//! `build_ai_tool_registry` 在 `AppState::init` 内、AppState 构建中被调用(state.rs),
|
||||
//! 此时无法传 `&AppState`(自引用 + 循环依赖,对齐 resolver 持 Arc<Database> 非 AppState 的
|
||||
//! 既有设计)。故把所需 Arc 句柄(db + agent 配置原子量 + LlmConcurrency)打包成
|
||||
//! [`GetAppConfigCtx`],state.rs init 先建这些 Arc(不依赖 db)再 build registry 注入。
|
||||
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize};
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use df_storage::crud::{AiProviderRepo, SettingsRepo};
|
||||
|
||||
use crate::state::LlmConcurrency;
|
||||
|
||||
/// 复用 IPC 层脱敏规则(单点维护,与 ai_list_providers 同口径)。
|
||||
use crate::commands::ai::commands::mask_api_key;
|
||||
/// keyring 异步解析(Tauri 单线程 runtime,同步 keyring 调用会卡)。
|
||||
use crate::commands::ai::secret::get_provider_secret_async;
|
||||
|
||||
/// app_settings KV 白名单(仅这些 key 进 LLM 视图,其余敏感/无关 KV 一律不返)。
|
||||
///
|
||||
/// 选取标准:与 AI 行为直接相关 + 无敏感凭证内容。`custom_prompt` 只返长度不返原文
|
||||
/// (用户可能写入含指令的私有提示,原文回灌进 LLM 上下文有指令注入风险 + 撑爆 prompt)。
|
||||
const ALLOWED_SETTING_KEYS: &[&str] = &[
|
||||
// 知识库行为配置(摘要/提炼开关,JSON 结构;值含开关不含凭证)
|
||||
"df-knowledge-config",
|
||||
// 审批超时(ms 数值,持久化版,与 AppState.approval_timeout_minutes 互补)
|
||||
"df-approval-timeout",
|
||||
];
|
||||
|
||||
/// get_app_config 工具注册期捕获的句柄集合(绕 AppState 自引用)。
|
||||
///
|
||||
/// 各字段均为 `Arc`(或内部全 Arc 的 Clone 廉价类型),clone 进工具闭包后与 AppState
|
||||
/// 字段共享同一底层原子量/Semaphore,故 handler 读到的永远是「当前生效值」
|
||||
/// (经 ai_set_concurrency_config / ai_set_agent_max_iterations 等 IPC 热改后立即反映)。
|
||||
///
|
||||
/// 字段对应 AppState 同名字段,state.rs init 先建这些 Arc 再组装 state(顺序调整,
|
||||
/// 因 build_ai_tool_registry 在 AppState 构建中被调,无法传 &AppState)。
|
||||
///
|
||||
/// `Clone` 廉价(全 Arc 字段):declare_tool! 宏把 ctx clone 进工具闭包,每次工具执行
|
||||
/// 再 clone 一份进 async move 块(宏展开语义,见 crates/df-ai/src/ai_tools_decl.rs)。
|
||||
#[derive(Clone)]
|
||||
pub struct GetAppConfigCtx {
|
||||
/// 数据库句柄(读 ai_providers / app_settings)
|
||||
pub db: Arc<df_storage::db::Database>,
|
||||
/// Agentic 循环最大轮次(热改即生效)
|
||||
pub agent_max_iterations: Arc<AtomicUsize>,
|
||||
/// 流式对话失败自动重试次数(热改即生效)
|
||||
pub agent_max_retries: Arc<AtomicUsize>,
|
||||
/// 审批超时分钟数(热改即生效)
|
||||
pub approval_timeout_minutes: Arc<AtomicU64>,
|
||||
/// LLM 并发控制(内部全 Arc,Clone 廉价,共享同一组 Semaphore)
|
||||
pub llm_concurrency: LlmConcurrency,
|
||||
}
|
||||
|
||||
/// 单个 provider 的脱敏视图(无 api_key 字段,绝不向 LLM 暴露任何 key 痕迹)。
|
||||
///
|
||||
/// 列表用:返所有 provider 概要(name/type/url/model/是否默认/是否启用),AI 知道有哪些
|
||||
/// provider 可用,不返 model_configs(详组能力配置,体积大且对「查当前配置」无必要)。
|
||||
fn provider_summary(p: &df_storage::models::AiProviderRecord) -> Value {
|
||||
json!({
|
||||
"name": p.name,
|
||||
"provider_type": p.provider_type,
|
||||
"base_url": p.base_url,
|
||||
"default_model": p.default_model,
|
||||
"is_default": p.is_default,
|
||||
"enabled": p.enabled,
|
||||
})
|
||||
}
|
||||
|
||||
/// 默认 provider 的扩展视图(含 api_key_masked,供 AI 判断 key 是否配置/有效)。
|
||||
///
|
||||
/// 仅默认这一条返脱敏 key(其他 provider 列表项不带 key 字段),平衡「AI 需知道 key 在不在」
|
||||
/// 与「最小披露」。空 key → 空串(明确告知 key 缺失,而非误导性地返 mask 占位)。
|
||||
fn default_provider_view(
|
||||
p: &df_storage::models::AiProviderRecord,
|
||||
real_key: &str,
|
||||
) -> Value {
|
||||
let mut view = provider_summary(p);
|
||||
// 作为对象才能插入字段(provider_summary 返的是 Object)。
|
||||
if let Some(obj) = view.as_object_mut() {
|
||||
let masked = if real_key.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
mask_api_key(real_key)
|
||||
};
|
||||
obj.insert("api_key_masked".into(), json!(masked));
|
||||
}
|
||||
view
|
||||
}
|
||||
|
||||
/// get_app_config 工具 handler 入口。
|
||||
///
|
||||
/// 返回结构:
|
||||
/// ```jsonc
|
||||
/// {
|
||||
/// "default_provider": { name, provider_type, base_url, default_model, is_default, enabled, api_key_masked: "sk-l••••3Yab" } | null,
|
||||
/// "providers": [{ name, provider_type, base_url, default_model, is_default, enabled }, ...],
|
||||
/// "agent": { max_iterations, max_retries, approval_timeout_minutes, concurrency: { per_conv } },
|
||||
/// "app_settings": { "<allowed_key>": "<value>", "custom_prompt_present": bool, "custom_prompt_length"?: number }
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// 全程只读:`list_all` / `load` / `get_all`,无任何 mutate 调用,无 IO 副作用。
|
||||
pub(crate) async fn execute_get_app_config(ctx: &GetAppConfigCtx) -> anyhow::Result<Value> {
|
||||
let providers_repo = AiProviderRepo::new(&ctx.db);
|
||||
let settings_repo = SettingsRepo::new(&ctx.db);
|
||||
|
||||
// ── providers:全部列表 + 默认项(带脱敏 key)──
|
||||
let providers = providers_repo
|
||||
.list_all()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("读 ai_providers 失败: {}", e))?;
|
||||
|
||||
// 默认 provider:取首条 is_default=true(list_all 已按 created_at DESC 排序,首条即最新)。
|
||||
// 不做 is_default 收敛写库(只读工具,收敛是 ai_list_providers IPC 的职责),脏数据照原样读。
|
||||
let default_idx = providers.iter().position(|p| p.is_default);
|
||||
|
||||
// 解析默认 provider 的真实 key(keyring 异步,仅这一条调 keyring,避免列表全量解析拖慢)。
|
||||
let default_real_key = match default_idx {
|
||||
Some(i) => {
|
||||
let p = &providers[i];
|
||||
if !p.api_key.is_empty() {
|
||||
p.api_key.clone() // 未迁移(老明文)
|
||||
} else {
|
||||
get_provider_secret_async(p.id.clone())
|
||||
.await
|
||||
.unwrap_or_default() // 迁移后从 keyring
|
||||
}
|
||||
}
|
||||
None => String::new(),
|
||||
};
|
||||
|
||||
let default_provider = match default_idx {
|
||||
Some(i) => json!(default_provider_view(&providers[i], &default_real_key)),
|
||||
None => Value::Null,
|
||||
};
|
||||
|
||||
// 列表:全部 provider 概要(无 key 字段)。
|
||||
let providers_view: Vec<Value> = providers.iter().map(provider_summary).collect();
|
||||
|
||||
// ── agent:内存原子量当前生效值(热改即反映,无需重启)──
|
||||
let max_iterations = ctx.agent_max_iterations.load(Ordering::SeqCst);
|
||||
let max_retries = ctx.agent_max_retries.load(Ordering::SeqCst);
|
||||
let approval_timeout_minutes = ctx.approval_timeout_minutes.load(Ordering::SeqCst);
|
||||
|
||||
// concurrency:per_conv permits 可读(AtomicUsize);global 无 getter(Semaphore permits
|
||||
// 封装在 Arc<Mutex<Arc<Semaphore>>> 内层无读取接口),只返 per_conv 当前生效值。
|
||||
// global 默认 3,经 ai_set_concurrency_config 热改后内存生效但不暴露读取,故此处省略
|
||||
// 不返猜测值,避免误导 LLM(若 LLM 需 global 当前值,后续可在 LlmConcurrency 补 getter)。
|
||||
let per_conv_permits = ctx.llm_concurrency.current_per_conv_permits();
|
||||
|
||||
let agent = json!({
|
||||
"max_iterations": max_iterations,
|
||||
"max_retries": max_retries,
|
||||
"approval_timeout_minutes": approval_timeout_minutes,
|
||||
"concurrency": {
|
||||
"per_conv": per_conv_permits,
|
||||
},
|
||||
});
|
||||
|
||||
// ── app_settings:白名单筛选(防敏感 KV 泄露 + 防 prompt 噪音)──
|
||||
let all_settings = settings_repo
|
||||
.get_all()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("读 app_settings 失败: {}", e))?;
|
||||
let mut app_settings = serde_json::Map::new();
|
||||
for (key, value) in all_settings {
|
||||
if ALLOWED_SETTING_KEYS.contains(&key.as_str()) {
|
||||
app_settings.insert(key, json!(value));
|
||||
}
|
||||
}
|
||||
// custom_prompt 仅返长度 + 是否存在,不返原文(防指令注入 + 防撑 prompt)。
|
||||
match settings_repo.get("custom_prompt").await {
|
||||
Ok(Some(cp)) => {
|
||||
let len = cp.chars().count();
|
||||
app_settings.insert("custom_prompt_present".into(), json!(true));
|
||||
app_settings.insert("custom_prompt_length".into(), json!(len));
|
||||
}
|
||||
_ => {
|
||||
app_settings.insert("custom_prompt_present".into(), json!(false));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(json!({
|
||||
"default_provider": default_provider,
|
||||
"providers": providers_view,
|
||||
"agent": agent,
|
||||
"app_settings": app_settings,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use df_storage::crud::AiProviderRepo;
|
||||
use df_storage::db::Database;
|
||||
use df_storage::models::AiProviderRecord;
|
||||
|
||||
/// 构造测试用 GetAppConfigCtx(默认值,handler 实际执行读真实 db 数据)。
|
||||
fn make_ctx(db: Arc<Database>) -> GetAppConfigCtx {
|
||||
GetAppConfigCtx {
|
||||
db,
|
||||
agent_max_iterations: Arc::new(AtomicUsize::new(10)),
|
||||
agent_max_retries: Arc::new(AtomicUsize::new(3)),
|
||||
approval_timeout_minutes: Arc::new(AtomicU64::new(15)),
|
||||
llm_concurrency: LlmConcurrency::new(3, 2),
|
||||
}
|
||||
}
|
||||
|
||||
/// 插入一条测试 provider(老明文路径:DB api_key 非空,不触发 keyring,测试无 OS 副作用)。
|
||||
async fn insert_provider(
|
||||
db: &Arc<Database>,
|
||||
id: &str,
|
||||
name: &str,
|
||||
api_key: &str,
|
||||
is_default: bool,
|
||||
enabled: bool,
|
||||
) {
|
||||
let repo = AiProviderRepo::new(db);
|
||||
repo.insert(AiProviderRecord {
|
||||
id: id.to_string(),
|
||||
name: name.to_string(),
|
||||
provider_type: "openai_compat".to_string(),
|
||||
api_key: api_key.to_string(),
|
||||
base_url: "https://api.example.com".to_string(),
|
||||
default_model: "glm-4-flash".to_string(),
|
||||
models: None,
|
||||
model_configs: Vec::new(),
|
||||
is_default,
|
||||
config: None,
|
||||
created_at: "0".to_string(),
|
||||
updated_at: "0".to_string(),
|
||||
enabled,
|
||||
weight: 50,
|
||||
})
|
||||
.await
|
||||
.expect("insert provider");
|
||||
}
|
||||
|
||||
/// 返回结构完整:四顶层键齐全 + providers 是数组。
|
||||
#[tokio::test]
|
||||
async fn returns_full_structure() {
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
insert_provider(&db, "p1", "默认", "sk-abcdef123456", true, true).await;
|
||||
let ctx = make_ctx(db);
|
||||
let res = execute_get_app_config(&ctx).await.unwrap();
|
||||
|
||||
// 四顶层键齐全
|
||||
assert!(res.get("default_provider").is_some(), "缺 default_provider");
|
||||
assert!(res.get("providers").is_some(), "缺 providers");
|
||||
assert!(res.get("agent").is_some(), "缺 agent");
|
||||
assert!(res.get("app_settings").is_some(), "缺 app_settings");
|
||||
// providers 是数组
|
||||
assert!(res["providers"].is_array(), "providers 应为数组");
|
||||
// agent 含四字段
|
||||
assert_eq!(res["agent"]["max_iterations"], 10);
|
||||
assert_eq!(res["agent"]["max_retries"], 3);
|
||||
assert_eq!(res["agent"]["approval_timeout_minutes"], 15);
|
||||
assert_eq!(res["agent"]["concurrency"]["per_conv"], 2);
|
||||
}
|
||||
|
||||
/// api_key 脱敏:默认 provider 返前 4 + •••• + 后 4,非空 key 不返明文。
|
||||
#[tokio::test]
|
||||
async fn default_provider_api_key_masked() {
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
// 13 字符 key:前 4 = sk-a, 后 4 = 3456, 中间 ••••
|
||||
insert_provider(&db, "p1", "默认", "sk-abcdef123456", true, true).await;
|
||||
let ctx = make_ctx(db);
|
||||
let res = execute_get_app_config(&ctx).await.unwrap();
|
||||
|
||||
let masked = res["default_provider"]["api_key_masked"]
|
||||
.as_str()
|
||||
.expect("default_provider 应含 api_key_masked");
|
||||
assert_eq!(masked, "sk-a••••3456", "13 字符 key 应脱敏为 前4+••••+后4");
|
||||
// 绝不含完整明文
|
||||
assert!(!masked.contains("abcdef"), "脱敏值绝不含明文中段");
|
||||
assert!(!masked.contains("1234"), "脱敏值绝不含明文尾段(除最后 4)");
|
||||
}
|
||||
|
||||
/// 列表项 providers 不含 api_key 字段(最小披露:仅默认项带脱敏 key)。
|
||||
#[tokio::test]
|
||||
async fn providers_list_has_no_api_key_field() {
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
insert_provider(&db, "p1", "默认", "sk-abcdef123456", true, true).await;
|
||||
insert_provider(&db, "p2", "备用", "sk-xxxxxxxxxxxx", false, true).await;
|
||||
let ctx = make_ctx(db);
|
||||
let res = execute_get_app_config(&ctx).await.unwrap();
|
||||
|
||||
let providers = res["providers"].as_array().unwrap();
|
||||
assert_eq!(providers.len(), 2, "应有 2 个 provider");
|
||||
for p in providers {
|
||||
assert!(
|
||||
p.get("api_key_masked").is_none(),
|
||||
"providers 列表项不应含 api_key_masked(仅默认项带)"
|
||||
);
|
||||
assert!(
|
||||
p.get("api_key").is_none(),
|
||||
"providers 列表项绝不含 api_key 字段"
|
||||
);
|
||||
// 列表项含必要概要字段
|
||||
assert!(p.get("name").is_some());
|
||||
assert!(p.get("provider_type").is_some());
|
||||
assert!(p.get("base_url").is_some());
|
||||
assert!(p.get("default_model").is_some());
|
||||
assert!(p.get("is_default").is_some());
|
||||
assert!(p.get("enabled").is_some());
|
||||
}
|
||||
}
|
||||
|
||||
/// 空 key(未配置)→ api_key_masked 为空串,明确告知缺失而非误导性占位。
|
||||
#[tokio::test]
|
||||
async fn empty_api_key_returns_empty_string() {
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
// DB api_key 空且 keyring 无此 id → resolve 返空 → handler 返空串
|
||||
// (keyring get_provider_secret_async 对不存在 id 返 None → unwrap_or_default 空串)
|
||||
insert_provider(&db, "p1", "无key", "", true, true).await;
|
||||
let ctx = make_ctx(db);
|
||||
let res = execute_get_app_config(&ctx).await.unwrap();
|
||||
|
||||
let masked = res["default_provider"]["api_key_masked"]
|
||||
.as_str()
|
||||
.expect("default_provider 应含 api_key_masked");
|
||||
assert_eq!(masked, "", "空 key 应返空串(明确缺失),非误导性 mask 占位");
|
||||
}
|
||||
|
||||
/// 无默认 provider → default_provider 为 null(不 panic,不臆造)。
|
||||
#[tokio::test]
|
||||
async fn no_default_provider_returns_null() {
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
insert_provider(&db, "p1", "非默认", "sk-abcdef123456", false, true).await;
|
||||
let ctx = make_ctx(db);
|
||||
let res = execute_get_app_config(&ctx).await.unwrap();
|
||||
|
||||
assert!(
|
||||
res["default_provider"].is_null(),
|
||||
"无默认 provider 时 default_provider 应为 null"
|
||||
);
|
||||
// providers 列表仍正常返
|
||||
assert_eq!(res["providers"].as_array().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
/// 短 key(≤8 字符)全脱敏为 •,不泄露长度外的信息(对齐 mask_api_key 规则)。
|
||||
#[tokio::test]
|
||||
async fn short_api_key_fully_masked() {
|
||||
let db = Arc::new(Database::open_in_memory().await.unwrap());
|
||||
insert_provider(&db, "p1", "短key", "sk-ab", true, true).await;
|
||||
let ctx = make_ctx(db);
|
||||
let res = execute_get_app_config(&ctx).await.unwrap();
|
||||
|
||||
let masked = res["default_provider"]["api_key_masked"]
|
||||
.as_str()
|
||||
.unwrap();
|
||||
assert_eq!(masked, "•••••", "5 字符 key 应全脱敏为 5 个 •(≤8 全脱敏)");
|
||||
}
|
||||
}
|
||||
@@ -13,7 +13,9 @@
|
||||
//! 4. **重定向限制 ≤3 跳**:每跳重新校验目标 host/IP(防 302 绕过到内网)。reqwest 默认跟 ≤10,
|
||||
//! 改用手动 redirect(Policy::none) + 自管循环。
|
||||
//! 5. **超时**:默认 30s,clamp ≤60s(对齐 run_command 思路,但封顶更紧——HTTP 调用应快速返回)。
|
||||
//! 6. **响应截断**:body ≤50KB(对齐 T-05),超长截断尾部保留 + truncated 标记。
|
||||
//! 6. **响应截断**:body 默认 ≤50KB(对齐 T-05),超长截断尾部保留 + truncated 标记。
|
||||
//! 截断阈值可经 `max_response_chars` 参数覆盖(clamp [1KB, 10MB]);或经 `output_file`
|
||||
//! 将响应全量落盘(不截断),body 字段置 null + path/bytes_written 返回。
|
||||
//!
|
||||
//! ## 风险分级
|
||||
//!
|
||||
@@ -26,12 +28,22 @@
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::net::IpAddr;
|
||||
use std::path::Path;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
/// 响应 body 截断上限(字节)。对齐 T-05 文档约定,防大响应撑爆 LLM context。
|
||||
const MAX_BODY_BYTES: usize = 50 * 1024;
|
||||
use crate::state::AllowedDirs;
|
||||
|
||||
/// 响应 body 截断默认上限(字节)。对齐 T-05 文档约定,防大响应撑爆 LLM context。
|
||||
///
|
||||
/// 2026-08-01:由硬编码改为默认值,LLM 可经 `max_response_chars` 参数覆盖(clamp [1024, 10MB])。
|
||||
/// 无 `max_response_chars` 时回退此默认(保旧行为零变更)。
|
||||
const DEFAULT_MAX_BODY_BYTES: usize = 50 * 1024;
|
||||
/// `max_response_chars` 参数 clamp 下限(字节)。防 LLM 传过小值截到无意义残片。
|
||||
const MIN_MAX_RESPONSE_CHARS: usize = 1024;
|
||||
/// `max_response_chars` 参数 clamp 上限(字节,10MB)。防 LLM 传超大值撑爆 LLM context。
|
||||
const MAX_MAX_RESPONSE_CHARS: usize = 10 * 1024 * 1024;
|
||||
|
||||
/// 默认请求超时(秒)。LLM 可经 args timeout_secs 覆盖。
|
||||
const DEFAULT_TIMEOUT_SECS: u64 = 30;
|
||||
@@ -246,9 +258,18 @@ fn build_request(
|
||||
/// - body: string,可选(GET 通常不传;POST/PUT/PATCH 传)
|
||||
/// - timeout_secs: 默认 30,clamp ≤60
|
||||
/// - parse: "json"|"text"|"auto"(默认 auto),控制 body 解析展示
|
||||
/// - output_file: string 可选,响应 body 写入指定文件路径(而非返回 JSON body),
|
||||
/// 适合大响应。路径走工具文件路径校验(validate_path + resolve_workspace_path_with_allowed
|
||||
/// + parent 授权目录校验,与 write_file 同源),全量写不截断。
|
||||
/// - max_response_chars: integer 可选,默认 50KB(51200),clamp [1024, 10MB]。
|
||||
/// body 截断阈值,替代硬编码上限。无 output_file 时控制返回 body 的截断;有 output_file 时不截断。
|
||||
///
|
||||
/// 返回 {status, status_text, headers, body, elapsed_ms, truncated, url(最终重定向后)}
|
||||
pub(crate) async fn execute_http_request(args: Value) -> anyhow::Result<Value> {
|
||||
/// `allowed`:授权目录快照(由 handler 闭包 read lock clone 传入),仅 output_file 落盘路径校验用。
|
||||
/// 无 output_file 时路径校验不触发,allowed 可为空快照(默认)。
|
||||
///
|
||||
/// 返回 {status, status_text, headers, body, elapsed_ms, truncated, url(最终重定向后)}。
|
||||
/// 有 output_file 时:body 字段置 null(已落盘),增 path/bytes_written,且 truncated=false(全量写)。
|
||||
pub(crate) async fn execute_http_request(args: Value, allowed: &AllowedDirs) -> anyhow::Result<Value> {
|
||||
// ── 参数解析 ──
|
||||
let method_str = args.get("method").and_then(|v| v.as_str()).unwrap_or("GET").to_uppercase();
|
||||
let method = match method_str.as_str() {
|
||||
@@ -286,6 +307,25 @@ pub(crate) async fn execute_http_request(args: Value) -> anyhow::Result<Value> {
|
||||
.min(MAX_TIMEOUT_SECS)
|
||||
.max(1); // clamp [1, 60]
|
||||
let parse_mode = args.get("parse").and_then(|v| v.as_str()).unwrap_or("auto").to_lowercase();
|
||||
// output_file:响应 body 落盘路径(可选)。存在时走文件路径校验 + 全量写,不返回 body。
|
||||
let output_file = args.get("output_file").and_then(|v| v.as_str()).map(|s| s.to_string());
|
||||
// max_response_chars:body 截断阈值(字节)。默认 DEFAULT_MAX_BODY_BYTES,clamp [1024, 10MB]。
|
||||
let max_response_chars = args.get("max_response_chars").and_then(|v| v.as_u64())
|
||||
.map(|n| (n as usize).clamp(MIN_MAX_RESPONSE_CHARS, MAX_MAX_RESPONSE_CHARS))
|
||||
.unwrap_or(DEFAULT_MAX_BODY_BYTES);
|
||||
// output_file 路径校验(若提供):validate_path + workspace 授权 + symlink 逃逸 +
|
||||
// parent 授权目录(落盘写需要,对齐 write_file FR-S8)。落盘时再次 resolve(复用同一校验)。
|
||||
if let Some(of) = &output_file {
|
||||
if of.trim().is_empty() {
|
||||
anyhow::bail!("output_file 不能为空");
|
||||
}
|
||||
let resolved = super::tool_registry::resolve_workspace_path_with_allowed(of, allowed)?;
|
||||
if let Some(parent) = Path::new(&resolved).parent() {
|
||||
if !allowed.is_authorized(parent) {
|
||||
anyhow::bail!("禁止在项目目录之外写入文件");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── SSRF 校验(三层:词法 URL → DNS resolve IP → 每跳重定向重复) ──
|
||||
let (scheme, host, port) = validate_url(&url_raw)?;
|
||||
@@ -314,11 +354,55 @@ pub(crate) async fn execute_http_request(args: Value) -> anyhow::Result<Value> {
|
||||
.or_insert(val);
|
||||
}
|
||||
|
||||
// body 读取 + 截断
|
||||
// body 读取
|
||||
let full_body = resp.bytes().await
|
||||
.map_err(|e| anyhow::anyhow!("读取响应 body 失败: {}", e))?;
|
||||
let total_bytes = full_body.len();
|
||||
let (body_text, truncated) = truncate_body(&full_body, MAX_BODY_BYTES);
|
||||
|
||||
// ── output_file 分支:全量写盘,不截断,不返回 body ──
|
||||
if let Some(of) = &output_file {
|
||||
let resolved = super::tool_registry::resolve_workspace_path_with_allowed(of, allowed)?;
|
||||
let path = resolved.to_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("output_file 路径含非法字符"))?
|
||||
.to_string();
|
||||
let target = Path::new(&path);
|
||||
// 父目录授权(对齐 write_file FR-S8)+ 自动创建(create_dir_all)
|
||||
if let Some(parent) = target.parent() {
|
||||
if !allowed.is_authorized(parent) {
|
||||
anyhow::bail!("禁止在项目目录之外写入文件");
|
||||
}
|
||||
tokio::fs::create_dir_all(parent).await
|
||||
.map_err(|e| anyhow::anyhow!("创建输出目录失败: {}", e))?;
|
||||
}
|
||||
// 全量写(原子:tmp→rename,避免写一半崩溃留半成品,与 write_file 同思路)
|
||||
let tmp = format!("{}.tmp-http", path);
|
||||
tokio::fs::write(&tmp, &full_body).await
|
||||
.map_err(|e| {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
anyhow::anyhow!("写入 output_file 失败: {}", e)
|
||||
})?;
|
||||
tokio::fs::rename(&tmp, &path).await
|
||||
.map_err(|e| {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
anyhow::anyhow!("output_file 原子重命名失败: {}", e)
|
||||
})?;
|
||||
return Ok(json!({
|
||||
"method": method_str,
|
||||
"url": final_url,
|
||||
"status": status,
|
||||
"status_text": status_text,
|
||||
"scheme": scheme,
|
||||
"headers": resp_headers,
|
||||
"body": null,
|
||||
"path": path,
|
||||
"bytes_written": total_bytes,
|
||||
"truncated": false,
|
||||
"elapsed_ms": elapsed_ms,
|
||||
}));
|
||||
}
|
||||
|
||||
// ── 默认分支:截断 + parse 格式化 + 返回 body ──
|
||||
let (body_text, truncated) = truncate_body(&full_body, max_response_chars);
|
||||
|
||||
// parse 处理:json → pretty;auto → 尝试 json 失败回退 text;二进制(含 \0)跳过
|
||||
let parsed_body = format_body(&body_text, total_bytes, &parse_mode, &resp_headers);
|
||||
@@ -594,14 +678,14 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn test_handler_missing_url_errors() {
|
||||
let args = json!({ "method": "GET" });
|
||||
let err = execute_http_request(args).await.unwrap_err();
|
||||
let err = execute_http_request(args, &AllowedDirs::default()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("缺少 url"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_invalid_method_errors() {
|
||||
let args = json!({ "method": "TRACE", "url": "https://example.com/" });
|
||||
let err = execute_http_request(args).await.unwrap_err();
|
||||
let err = execute_http_request(args, &AllowedDirs::default()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("不支持的 method"));
|
||||
}
|
||||
|
||||
@@ -609,21 +693,21 @@ mod tests {
|
||||
async fn test_handler_rejects_localhost_url() {
|
||||
// SSRF 防护:handler 入口 validate_url 拒 localhost,不发请求
|
||||
let args = json!({ "method": "GET", "url": "http://localhost:8080/admin" });
|
||||
let err = execute_http_request(args).await.unwrap_err();
|
||||
let err = execute_http_request(args, &AllowedDirs::default()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("localhost") || format!("{}", err).contains("SSRF"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_rejects_private_ip_url() {
|
||||
let args = json!({ "method": "GET", "url": "http://169.254.169.254/latest/meta-data/" });
|
||||
let err = execute_http_request(args).await.unwrap_err();
|
||||
let err = execute_http_request(args, &AllowedDirs::default()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("私网") || format!("{}", err).contains("SSRF"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handler_rejects_non_http_scheme() {
|
||||
let args = json!({ "url": "file:///etc/passwd" });
|
||||
let err = execute_http_request(args).await.unwrap_err();
|
||||
let err = execute_http_request(args, &AllowedDirs::default()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("协议"));
|
||||
}
|
||||
|
||||
@@ -632,7 +716,44 @@ mod tests {
|
||||
// timeout_secs=9999 应 clamp 到 60;此处只验证不因超大值 panic(handler 内会走网络,
|
||||
// 故本测试不期望成功完成——用 localhost 触发 SSRF 早期拒绝,验证 clamp 在拒绝前不 panic)
|
||||
let args = json!({ "url": "http://localhost/", "timeout_secs": 9999 });
|
||||
let _ = execute_http_request(args).await; // 不 panic 即通过
|
||||
let _ = execute_http_request(args, &AllowedDirs::default()).await; // 不 panic 即通过
|
||||
}
|
||||
|
||||
// ── output_file:路径校验拒绝(提前在 SSRF/请求前触发,不发网络) ──
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_output_file_rejects_path_traversal() {
|
||||
// output_file 含 .. → validate_path 拒绝(在发请求前,故无需真实网络)
|
||||
let args = json!({
|
||||
"url": "https://example.com/",
|
||||
"output_file": "../escape.txt"
|
||||
});
|
||||
let err = execute_http_request(args, &AllowedDirs::default()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("路径遍历"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_output_file_rejects_empty() {
|
||||
let args = json!({
|
||||
"url": "https://example.com/",
|
||||
"output_file": " "
|
||||
});
|
||||
let err = execute_http_request(args, &AllowedDirs::default()).await.unwrap_err();
|
||||
assert!(format!("{}", err).contains("output_file 不能为空"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_output_file_rejects_relative_without_allowed() {
|
||||
// 相对路径 + 空白名单 → resolve_workspace_path_with_allowed 引导报错(在发请求前)
|
||||
let args = json!({
|
||||
"url": "https://example.com/",
|
||||
"output_file": "out.txt"
|
||||
});
|
||||
let err = execute_http_request(args, &AllowedDirs::default()).await.unwrap_err();
|
||||
// 无授权目录时 resolve 报"请先绑定项目目录"
|
||||
let msg = format!("{}", err);
|
||||
assert!(msg.contains("绑定项目目录") || msg.contains("授权目录"),
|
||||
"期望授权引导错误,实际: {}", msg);
|
||||
}
|
||||
|
||||
// ── 真实网络集成(#[ignore]:CI 无网时跳过,本地手跑) ──
|
||||
@@ -641,7 +762,7 @@ mod tests {
|
||||
#[ignore = "需真实网络(example.com),CI 无网时跳过:cargo test -- --ignored"]
|
||||
async fn integration_get_example_com_200() {
|
||||
let args = json!({ "method": "GET", "url": "https://example.com/" });
|
||||
let result = execute_http_request(args).await.expect("GET example.com 应成功");
|
||||
let result = execute_http_request(args, &AllowedDirs::default()).await.expect("GET example.com 应成功");
|
||||
let status = result["status"].as_u64().unwrap();
|
||||
assert!(status == 200, "期望 200,实际 {}", status);
|
||||
assert!(result["body"].as_str().unwrap_or("").contains("Example Domain"));
|
||||
@@ -649,13 +770,58 @@ mod tests {
|
||||
assert!(result["elapsed_ms"].as_u64().unwrap_or(0) > 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "需真实网络,本地手跑"]
|
||||
async fn integration_max_response_chars_truncates() {
|
||||
// 默认 50KB 不截断 example.com(<1KB);max_response_chars=100 强制截断验证参数生效
|
||||
let args = json!({
|
||||
"url": "https://example.com/",
|
||||
"max_response_chars": 100
|
||||
});
|
||||
let result = execute_http_request(args, &AllowedDirs::default()).await
|
||||
.expect("GET example.com 应成功");
|
||||
assert_eq!(result["truncated"], true, "应被 max_response_chars 截断");
|
||||
// 默认(无 max_response_chars)不截断
|
||||
let args2 = json!({ "url": "https://example.com/" });
|
||||
let result2 = execute_http_request(args2, &AllowedDirs::default()).await.unwrap();
|
||||
assert_eq!(result2["truncated"], false);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "需真实网络 + 临时目录,本地手跑"]
|
||||
async fn integration_output_file_writes_body() {
|
||||
// output_file 落盘:body 应为 null,path/bytes_written 存在,truncated=false。
|
||||
// 路径校验需授权目录,故构建含 temp_dir 的 AllowedDirs(对齐 write_file 测试思路)。
|
||||
let dir = std::env::temp_dir();
|
||||
let out_path = dir.join("devflow_http_test_output.html");
|
||||
let out_str = out_path.to_string_lossy().to_string();
|
||||
let mut allowed = AllowedDirs::default();
|
||||
allowed.persistent.insert(dir.clone());
|
||||
let args = json!({
|
||||
"url": "https://example.com/",
|
||||
"output_file": out_str
|
||||
});
|
||||
let result = execute_http_request(args, &allowed).await
|
||||
.expect("GET example.com 应成功");
|
||||
assert_eq!(result["status"], 200);
|
||||
assert_eq!(result["body"], serde_json::Value::Null, "body 应置 null(已落盘)");
|
||||
assert_eq!(result["truncated"], false, "落盘不截断");
|
||||
let bytes_written = result["bytes_written"].as_u64().unwrap();
|
||||
assert!(bytes_written > 0, "应写入非零字节");
|
||||
assert_eq!(result["path"], out_str);
|
||||
// 校验文件内容确实是 example.com 响应
|
||||
let written = std::fs::read_to_string(&out_path).unwrap();
|
||||
assert!(written.contains("Example Domain"));
|
||||
let _ = std::fs::remove_file(&out_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "需真实网络,本地手跑"]
|
||||
async fn integration_get_follows_redirect() {
|
||||
// httpbin.org/redirect-to 重定向到 example.com,验证重定向链 + SSRF 每跳校验通过
|
||||
let target = urlencoding("https://example.com/");
|
||||
let args = json!({ "url": format!("https://httpbin.org/redirect-to?url={}", target) });
|
||||
let result = execute_http_request(args).await.expect("重定向链应成功");
|
||||
let result = execute_http_request(args, &AllowedDirs::default()).await.expect("重定向链应成功");
|
||||
let status = result["status"].as_u64().unwrap_or(0);
|
||||
// 最终落地 200(经重定向)
|
||||
assert!(status >= 200 && status < 400, "重定向后状态: {}", status);
|
||||
|
||||
@@ -37,6 +37,10 @@ pub mod conversation;
|
||||
pub mod event_bus;
|
||||
pub mod http;
|
||||
pub mod fetch_url;
|
||||
pub mod fetch_search;
|
||||
pub mod download_file;
|
||||
pub(crate) mod generate_image;
|
||||
pub mod get_app_config;
|
||||
pub mod knowledge_inject;
|
||||
pub mod prompt;
|
||||
pub mod provider_pool;
|
||||
|
||||
@@ -301,6 +301,52 @@ pub(crate) fn compute_file_hash(meta: &std::fs::Metadata) -> String {
|
||||
format!("{}_{}", modified.unwrap_or(0), meta.len())
|
||||
}
|
||||
|
||||
/// 将原始字节解码为 UTF-8 字符串，优先识别 UTF-16 LE/BE BOM 解码转 UTF-8。
|
||||
///
|
||||
/// 背景(P1 修复):PowerShell `Out-File`/`>` 重定向默认产 UTF-16 LE BOM 文件,
|
||||
/// 每个 ASCII 字符高字节为 0x00(NUL)。旧 read_file 用 read_to_string 直接读 →
|
||||
/// (a) read_to_string 对含 NUL 字节内容硬失败 InvalidData;或 (b) 即便侥幸读到,
|
||||
/// 后续 is_binary / patch_file 的 `\0` 检测把 UTF-16 文本误判二进制拒读 →
|
||||
/// run_command(Out-File) → 文件 → read_file 链路在 Windows 断裂。
|
||||
///
|
||||
/// 本 helper 替代 read_to_string 的纯字节→字符串转换:
|
||||
/// 1. UTF-16 LE BOM(FF FE):strip BOM 后 u16::from_le_bytes + String::from_utf16 解码;
|
||||
/// 2. UTF-16 BE BOM(FE FF):同上 from_be_bytes;
|
||||
/// 3. 无 BOM:走 String::from_utf8(等价原 read_to_string 的 UTF-8 校验)。
|
||||
/// 注:UTF-8 BOM(EF BB BF)是合法 UTF-8 序列,read_to_string 原本就能读(内容首字符多 U+FEFF),
|
||||
/// 不在本 helper 误判路径内,无需特判。
|
||||
///
|
||||
/// 解码失败(真二进制/非 UTF-8 非 BOM 内容)返 io::Error(InvalidData),与 read_to_string
|
||||
/// 同语义——调用方保留原 `if e.kind() == InvalidData { 返 binary 标记 }` 分支不变,
|
||||
/// 真二进制(图片/编译产物)仍被拦截,UTF-16 文本因 BOM 优先识别不再进 \x00 误判路径。
|
||||
///
|
||||
/// 奇数字节 UTF-16(尾字节落单):drop 最后一个残字节后解码(对齐 PowerShell 偶发写半字符场景,
|
||||
/// 不 bail 让多数可读内容通过)。Lone surrogate 交 String::from_utf16 lossy 兜底不 panic。
|
||||
pub(crate) fn decode_bytes_to_string(bytes: &[u8]) -> std::io::Result<String> {
|
||||
// UTF-16 LE BOM: FF FE
|
||||
if bytes.len() >= 2 && bytes[0] == 0xFF && bytes[1] == 0xFE {
|
||||
let body = &bytes[2..];
|
||||
let units: Vec<u16> = body
|
||||
.chunks_exact(2)
|
||||
.map(|c| u16::from_le_bytes([c[0], c[1]]))
|
||||
.collect();
|
||||
// from_utf16 遇 lone surrogate 返 Err(此处用 lossy 兜底保内容可读,不阻断)
|
||||
return Ok(String::from_utf16_lossy(&units));
|
||||
}
|
||||
// UTF-16 BE BOM: FE FF
|
||||
if bytes.len() >= 2 && bytes[0] == 0xFE && bytes[1] == 0xFF {
|
||||
let body = &bytes[2..];
|
||||
let units: Vec<u16> = body
|
||||
.chunks_exact(2)
|
||||
.map(|c| u16::from_be_bytes([c[0], c[1]]))
|
||||
.collect();
|
||||
return Ok(String::from_utf16_lossy(&units));
|
||||
}
|
||||
// 无 BOM:走 UTF-8 校验(等价 read_to_string,失败 InvalidData → 调用方判二进制)
|
||||
String::from_utf8(bytes.to_vec())
|
||||
.map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "非 UTF-8 / 非 UTF-16 BOM 文本"))
|
||||
}
|
||||
|
||||
/// 阶段4(容错/恢复,开关 `df-ai-approval-retry`):跨盘/跨卷文件移动统一降级 helper。
|
||||
///
|
||||
/// 背景:Windows 跨盘符(C→E)或跨卷时 `tokio::fs::rename` 报 `os error 17`
|
||||
@@ -448,16 +494,26 @@ pub(crate) async fn bind_dir_to_project(
|
||||
/// F-260619-03 Phase A: 新增 `allowed_dirs: &Arc<RwLock<AllowedDirs>>` 形参,
|
||||
/// 文件工具闭包 clone 进闭包,在 resolve_workspace_path_with_allowed 校验动态白名单。
|
||||
/// data/http 工具不涉及文件路径,不传白名单。
|
||||
///
|
||||
/// get_app_config(2026-08-02):新增 `get_app_config_ctx: GetAppConfigCtx` 形参,
|
||||
/// 该工具查 DevFlow 自身当前 AI 配置(provider/agent/settings),需读 AppState 内存原子量
|
||||
/// (agent_max_iterations 等)+ db。因本函数在 AppState::init 内、AppState 构建中被调,
|
||||
/// 无法传 &AppState(自引用),故把所需 Arc 句柄打包成 GetAppConfigCtx 传入(state.rs
|
||||
/// init 先建这些 Arc 再 build registry 再组装 state,顺序调整)。
|
||||
pub fn build_ai_tool_registry(
|
||||
db: &Arc<Database>,
|
||||
allowed_dirs: &Arc<RwLock<AllowedDirs>>,
|
||||
data_dir: PathBuf,
|
||||
get_app_config_ctx: super::get_app_config::GetAppConfigCtx,
|
||||
) -> AiToolRegistry {
|
||||
let mut registry = AiToolRegistry::new();
|
||||
register_data_tools(&mut registry, db);
|
||||
register_file_tools(&mut registry, allowed_dirs, data_dir);
|
||||
register_http_tools(&mut registry);
|
||||
register_http_tools(&mut registry, allowed_dirs);
|
||||
register_fetch_url_tool(&mut registry);
|
||||
register_fetch_search_tool(&mut registry);
|
||||
register_generate_image_tool(&mut registry, db, allowed_dirs);
|
||||
register_get_app_config_tool(&mut registry, get_app_config_ctx);
|
||||
registry
|
||||
}
|
||||
|
||||
@@ -473,8 +529,11 @@ pub fn build_ai_tool_registry(
|
||||
/// async move 块(逻辑零变更),仅闭包包装改由 `declare_tool!` 宏生成。本函数改 thin 委托
|
||||
/// super::tools::http::register(...)。基线测试 test_build_ai_tool_registry_baseline_tool_count
|
||||
/// 仍断言 48 总量 + 工具名集合稳定。
|
||||
fn register_http_tools(registry: &mut AiToolRegistry) {
|
||||
super::tools::http::register(registry);
|
||||
///
|
||||
/// 2026-08-01 output_file 参数:响应可落盘,路径走文件路径校验,故透传 allowed_dirs 给
|
||||
/// tools::http::register(由其闭包捕获,handler 内 read lock clone 传 execute_http_request)。
|
||||
fn register_http_tools(registry: &mut AiToolRegistry, allowed_dirs: &Arc<RwLock<AllowedDirs>>) {
|
||||
super::tools::http::register(registry, allowed_dirs);
|
||||
}
|
||||
|
||||
/// fetch_url AI 工具注册(1 个:URL → markdown 文档嗅探)— 只读 GET 网页文档。
|
||||
@@ -490,6 +549,63 @@ fn register_fetch_url_tool(registry: &mut AiToolRegistry) {
|
||||
super::tools::fetch_url::register(registry);
|
||||
}
|
||||
|
||||
/// fetch_search AI 工具注册(1 个:搜索引擎查询)— DuckDuckGo HTML 端点免 key 搜索。
|
||||
/// 不持 db,纯 reqwest GET + HTML 字符串扫描解析。SSRF 防护复用 commands/ai/http.rs
|
||||
/// (validate_url / resolve_and_check_host / build_client / execute_with_redirects)。
|
||||
///
|
||||
/// 风险:Low(只读 GET 搜索,无副作用,与 fetch_url 同级)。与 fetch_url 分工:fetch_url 输出
|
||||
/// 已知 URL 的 markdown 正文,fetch_search 输出未知关键词的候选 URL 列表(标题+摘要),
|
||||
/// 供 LLM 挑选入口再 fetch_url 深读。补「fetch_url 只抓已知 URL,无搜索能力」缺口。
|
||||
///
|
||||
/// HTML 解析:手写字符串扫描(对齐 fetch_url::extract_title 风格),不引 scraper 重依赖
|
||||
/// (html5ever+selectors+cssparser 编译产物大,DDG 单一来源不值)。基线测试守护总量 + 工具名集合。
|
||||
fn register_fetch_search_tool(registry: &mut AiToolRegistry) {
|
||||
super::tools::fetch_search::register(registry);
|
||||
}
|
||||
|
||||
/// generate_image AI 工具注册(1 个:调 provider 图像生成端点 + 下载落地)。
|
||||
///
|
||||
/// 工具职责:调 provider 的 OpenAI 兼容 `/v1/images/generations` 端点生成图片(默认 SenseNova
|
||||
/// U1 Fast,也支持 DALL-E 等),自动下载到工作区 output_path。图像生成模型非 chat,走独立
|
||||
/// 端点非 chat completions,故不能经主对话链路调用,必须用本工具显式触发。
|
||||
///
|
||||
/// 涉网络(provider POST + 图片 URL 下载)+ 付费 API(按张计费)+ 写文件(落盘 output_path),
|
||||
/// 与 download_file 同源(均持 allowed_dirs 做 output_path 白名单校验 + 默认路径锚定)。
|
||||
/// risk=High(付费 + 写磁盘 + 外发)。基线测试 test_build_ai_tool_registry_baseline_tool_count
|
||||
/// 守护总量 + 工具名集合稳定。
|
||||
///
|
||||
/// handler 在 commands/ai/generate_image.rs,声明式注册在 tools/generate_image.rs。
|
||||
/// declare_tool! 单捕获限制:tools 层把 (db, allowed_dirs) 包进 tuple Arc 在闭包内解构。
|
||||
fn register_generate_image_tool(
|
||||
registry: &mut AiToolRegistry,
|
||||
db: &Arc<Database>,
|
||||
allowed_dirs: &Arc<RwLock<AllowedDirs>>,
|
||||
) {
|
||||
super::tools::generate_image::register(registry, db, allowed_dirs);
|
||||
}
|
||||
|
||||
/// get_app_config AI 工具注册(1 个:查 DevFlow 自身当前 AI 配置,只读 + 脱敏)—
|
||||
/// 返 {default_provider, providers, agent, app_settings} JSON。
|
||||
///
|
||||
/// 治症状:AI 想看「当前 provider/model/agent 设置」时,因没原生查配置工具,只能 run_command
|
||||
/// 执行 PowerShell 内联 node/python 脚本查 db,引号三重嵌套必失败。本工具一调即得,根治绕行。
|
||||
///
|
||||
/// 风险:Low(纯只读 list/load,无写副作用)。安全:api_key 经 mask_api_key 脱敏
|
||||
/// (前 4 + •••• + 后 4),绝不返明文;app_settings 走白名单不返敏感 KV。
|
||||
///
|
||||
/// handler 在 commands/ai/get_app_config.rs,声明式注册在 tools/get_app_config.rs。
|
||||
/// declare_tool! 单捕获:把 db + 3 个 agent 配置 Arc + LlmConcurrency 包进 GetAppConfigCtx
|
||||
/// struct(Clone 廉价,全 Arc 字段),闭包内 &ctx 转调 handler。
|
||||
///
|
||||
/// ctx 字段与 AppState 同名字段共享同一底层原子量(经 ai_set_concurrency_config /
|
||||
/// ai_set_agent_max_iterations 等 IPC 热改后,handler 读到的永远是当前生效值)。
|
||||
fn register_get_app_config_tool(
|
||||
registry: &mut AiToolRegistry,
|
||||
ctx: super::get_app_config::GetAppConfigCtx,
|
||||
) {
|
||||
super::tools::get_app_config::register(registry, ctx);
|
||||
}
|
||||
|
||||
/// 数据层 AI 工具注册(25 个持 db 的 CRUD/状态机/工作流/知识图谱工具)——从 build_ai_tool_registry 抽出。
|
||||
///
|
||||
/// 工具闭包捕获 `db: &Arc<Database>` Arc 重建 Repo(列表/创建/更新/删除/状态推进/工作流/任务关联)。
|
||||
@@ -639,6 +755,10 @@ fn register_file_tools(
|
||||
data_dir: PathBuf,
|
||||
) {
|
||||
super::tools::file::register(registry, allowed_dirs, data_dir);
|
||||
// download_file:跨平台 URL→文件流式下载(reqwest bytes_stream + tokio::fs)。
|
||||
// 写文件需路径授权(与 write_file 同源),故在 register_file_tools 注册(已有 allowed_dirs)。
|
||||
// SSRF 复用 commands/ai/http.rs,大小上限 200MB + 原子写(tmp→rename)。
|
||||
super::tools::download_file::register(registry, allowed_dirs);
|
||||
}
|
||||
|
||||
/// 探测可执行文件路径(which/where 风格),8s 超时兜底。
|
||||
@@ -1122,6 +1242,28 @@ pub(crate) fn search_files_recursive<'a>(
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
// get_app_config baseline 测试辅助:构造 GetAppConfigCtx(与 AppState init 同款默认值,
|
||||
// 仅用于 build_ai_tool_registry 注册阶段定义完整性断言,handler 不实际执行故无需真实配置)。
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize};
|
||||
use crate::state::LlmConcurrency;
|
||||
|
||||
/// 测试用 GetAppConfigCtx 构造器(默认值对齐 AppState::init:iter=10/retries=3/timeout=15/concurrency 3,2)。
|
||||
///
|
||||
/// build_ai_tool_registry 第 4 参,5 处测试调用共用,避免每处内联 5 行重复。
|
||||
/// Arc 句柄独立于 AppState(测试无 AppState),仅满足注册期类型签名,handler 不执行。
|
||||
fn make_test_get_app_config_ctx(db: &Arc<Database>) -> super::get_app_config::GetAppConfigCtx {
|
||||
super::get_app_config::GetAppConfigCtx {
|
||||
db: db.clone(),
|
||||
agent_max_iterations: Arc::new(AtomicUsize::new(
|
||||
crate::commands::ai::agentic::DEFAULT_MAX_AGENT_ITERATIONS,
|
||||
)),
|
||||
agent_max_retries: Arc::new(AtomicUsize::new(
|
||||
crate::commands::ai::agentic::DEFAULT_MAX_AGENT_RETRIES,
|
||||
)),
|
||||
approval_timeout_minutes: Arc::new(AtomicU64::new(15)),
|
||||
llm_concurrency: LlmConcurrency::new(3, 2),
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// 工具注册基线测试（SMELL-P0-2 拆分防护）
|
||||
@@ -1131,7 +1273,7 @@ mod tests {
|
||||
// 任一层漏移 register 调用,此测试立即红。工具名集合也断言,防 rename 致 LLM tool 突变。
|
||||
// ============================================================
|
||||
|
||||
/// build_ai_tool_registry 应注册恰好 51 个工具(36 data + 13 file + 1 http + 1 fetch_url),且工具名集合稳定。
|
||||
/// build_ai_tool_registry 应注册恰好 55 个工具(36 data + 14 file + 1 http + 1 fetch_url + 1 fetch_search + 1 generate_image + 1 get_app_config),且工具名集合稳定。
|
||||
///
|
||||
/// 用 in-memory SQLite(Database::open_in_memory 自跑迁移),构造零外部依赖的 db,
|
||||
// 不实际执行任何 handler——仅断言注册阶段的定义完整性,故无需真实数据。
|
||||
@@ -1142,7 +1284,7 @@ mod tests {
|
||||
// F-260619-03 Phase A: build_ai_tool_registry 新增 allowed_dirs 形参,
|
||||
// 测试用 default_with_root(仅 workspace_root),零回归(白名单含 workspace_root)。
|
||||
let allowed_dirs = Arc::new(RwLock::new(AllowedDirs::default_with_root()));
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""));
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""), make_test_get_app_config_ctx(&db));
|
||||
|
||||
// 总量基线:41(27 data + 13 file + 1 http)。拆分前后必须一致。
|
||||
// F-260621: file 层 10→11(新增 grep 跨文件内容搜索工具)。
|
||||
@@ -1168,15 +1310,31 @@ mod tests {
|
||||
// fetch_url(2026-08-01): http 层 1→2(新增 URL→markdown 文档嗅探,只读 GET,
|
||||
// htmd HTML→markdown + 去噪音 + 截断,SSRF 复用 http.rs)。
|
||||
// 51 = 36 data + 13 file + 1 http + 1 fetch_url。
|
||||
// download_file(2026-08-01): file 层 13→14(新增 URL→文件跨平台流式下载,
|
||||
// reqwest bytes_stream + tokio::fs,替代 curl/wget alias 陷阱。写文件需 allowed_dirs
|
||||
// 故注册在 register_file_tools;SSRF 复用 http.rs + 路径校验复用 validate_path)。
|
||||
// 52 = 36 data + 14 file + 1 http + 1 fetch_url。
|
||||
// fetch_search(2026-08-02): fetch_search 层 1→2(新增搜索引擎查询,DuckDuckGo HTML 免 key,
|
||||
// GET + HTML 字符串扫描解析结果列表,补 fetch_url 只抓已知 URL 缺口。SSRF 复用 http.rs)。
|
||||
// 53 = 36 data + 14 file + 1 http + 1 fetch_url + 1 fetch_search。
|
||||
// generate_image(2026-08-02): generate_image 层 1→1(新增调 provider 图像生成端点 +
|
||||
// 下载落地,默认 sensenova-u1-fast,也支持 DALL-E。SSRF 复用 http.rs + 路径校验复用
|
||||
// resolve_workspace_path_with_allowed,declare_tool! 单捕获绕法用 tuple Arc)。
|
||||
// 54 = 36 data + 14 file + 1 http + 1 fetch_url + 1 fetch_search + 1 generate_image。
|
||||
// get_app_config(2026-08-02): get_app_config 层 1→1(新增查 DevFlow 自身当前 AI 配置,
|
||||
// 只读 + api_key 脱敏,治 AI 查配置绕 run_command PowerShell 内联脚本引号嵌套失败。
|
||||
// build_ai_tool_registry 加第 4 参 GetAppConfigCtx 绕 AppState 自引用,declare_tool!
|
||||
// 单捕获用 struct ctx 包 db + 3 agent Arc + LlmConcurrency,共享 AppState 同名句柄)。
|
||||
// 55 = 36 data + 14 file + 1 http + 1 fetch_url + 1 fetch_search + 1 generate_image + 1 get_app_config。
|
||||
assert_eq!(
|
||||
registry.len(),
|
||||
51,
|
||||
"工具总数应为 51(36 data + 13 file + 1 http + 1 fetch_url),实际 {}", registry.len()
|
||||
55,
|
||||
"工具总数应为 55(36 data + 14 file + 1 http + 1 fetch_url + 1 fetch_search + 1 generate_image + 1 get_app_config),实际 {}", registry.len()
|
||||
);
|
||||
|
||||
// 工具名集合基线:防 rename / 漏注册 / 误删除。
|
||||
// data 层 36 个(持 db):CRUD/状态机/工作流/知识图谱任务关联 + 项目事件流 + 基础设施配置 + git 工具
|
||||
// file 层 13 个(不持 db):命令/读/列/写/改/元/追加/删/移/搜/grep/环境探测/符号解析
|
||||
// file 层 14 个(不持 db):命令/读/列/写/改/元/追加/删/移/搜/grep/环境探测/符号解析/下载
|
||||
// http 层 1 个(不持 db):http_request
|
||||
let mut expected: Vec<&str> = vec![
|
||||
// ── data 层 (36) ──
|
||||
@@ -1200,17 +1358,27 @@ mod tests {
|
||||
"git_status", "git_diff", "git_log",
|
||||
// Git 写工具
|
||||
"git_commit", "git_branch", "git_merge",
|
||||
// ── file 层 (13) ──（run_command 注册顺序已移至末位降低 LLM 偏好,
|
||||
// ── file 层 (14) ──（run_command 注册顺序已移至末位降低 LLM 偏好,
|
||||
// 集合断言经 sort 后与顺序无关,仅守护工具名不漂移。grep 新增 F-260621;
|
||||
// detect_environment 新增 L1 环境感知 设计 §2.1;read_symbol 新增 AST 代码智能）
|
||||
// detect_environment 新增 L1 环境感知 设计 §2.1;read_symbol 新增 AST 代码智能;
|
||||
// download_file 新增跨平台 URL→文件流式下载）
|
||||
"read_file", "read_symbol", "list_directory", "write_file",
|
||||
"patch_file", "file_info", "append_file",
|
||||
"delete_file", "rename_file", "search_files", "run_command",
|
||||
"grep", "detect_environment",
|
||||
// download_file(URL→文件流式下载,写文件需 allowed_dirs,注册在 register_file_tools)
|
||||
"download_file",
|
||||
// ── http 层 (1) ──
|
||||
"http_request",
|
||||
// ── fetch_url 层 (1) ──(URL → markdown 文档嗅探,只读 GET,与 http_request 分工)
|
||||
"fetch_url",
|
||||
// ── fetch_search 层 (1) ──(搜索引擎查询 DDG HTML 免 key,只读 GET,补 fetch_url 缺口)
|
||||
"fetch_search",
|
||||
// ── generate_image 层 (1) ──(调 provider 图像生成端点 + 下载落地,默认 sensenova-u1-fast)
|
||||
"generate_image",
|
||||
// ── get_app_config 层 (1) ──(查 DevFlow 自身当前 AI 配置,只读 + api_key 脱敏,
|
||||
// 治 AI 查配置绕 run_command PowerShell 内联脚本引号嵌套失败)
|
||||
"get_app_config",
|
||||
];
|
||||
expected.sort_unstable();
|
||||
|
||||
@@ -1412,7 +1580,7 @@ mod tests {
|
||||
|
||||
let db = Database::open_in_memory().await.expect("in-memory db 初始化失败");
|
||||
let db = Arc::new(db);
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""));
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""), make_test_get_app_config_ctx(&db));
|
||||
let canon_file = file.canonicalize().unwrap().to_string_lossy().to_string();
|
||||
let args = serde_json::json!({ "path": canon_file, "limit": 15 });
|
||||
let res = registry.execute("read_file", args).await.expect("read_file 执行失败");
|
||||
@@ -1439,7 +1607,7 @@ mod tests {
|
||||
|
||||
let db = Database::open_in_memory().await.expect("in-memory db 初始化失败");
|
||||
let db = Arc::new(db);
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""));
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""), make_test_get_app_config_ctx(&db));
|
||||
let canon_file = file.canonicalize().unwrap().to_string_lossy().to_string();
|
||||
let args = serde_json::json!({ "path": canon_file });
|
||||
let res = registry.execute("read_file", args).await.expect("read_file 执行失败");
|
||||
@@ -1450,6 +1618,114 @@ mod tests {
|
||||
fs::remove_dir_all(&tmp).ok();
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// P1 修复:UTF-16 LE/BE BOM 文件解码(PowerShell Out-File/> 产 UTF-16 LE,
|
||||
// 旧 read_to_string 因 ASCII 高字节 0x00 判 InvalidData → 误判二进制拒读)
|
||||
// decode_bytes_to_string 纯函数单测 + read_file 端到端 + 真 binary 仍拒读
|
||||
// ============================================================
|
||||
|
||||
/// decode_bytes_to_string: UTF-16 LE BOM(FF FE) → 可读 UTF-8
|
||||
#[test]
|
||||
fn test_decode_bytes_utf16_le_bom() {
|
||||
// "hello\n" UTF-16 LE 编码
|
||||
let text = "hello\n";
|
||||
let mut bytes: Vec<u8> = vec![0xFF, 0xFE]; // LE BOM
|
||||
for u in text.encode_utf16() {
|
||||
bytes.extend_from_slice(&u.to_le_bytes());
|
||||
}
|
||||
let s = decode_bytes_to_string(&bytes).expect("UTF-16 LE BOM 应解码成功");
|
||||
assert_eq!(s, text, "UTF-16 LE BOM 文件应解码为原文本");
|
||||
}
|
||||
|
||||
/// decode_bytes_to_string: UTF-16 BE BOM(FE FF) → 可读 UTF-8
|
||||
#[test]
|
||||
fn test_decode_bytes_utf16_be_bom() {
|
||||
// 中文 + ASCII,UTF-16 BE 编码(验证非 ASCII 码点 + 大端序)
|
||||
let text = "中文 ascii\n";
|
||||
let mut bytes: Vec<u8> = vec![0xFE, 0xFF]; // BE BOM
|
||||
for u in text.encode_utf16() {
|
||||
bytes.extend_from_slice(&u.to_be_bytes());
|
||||
}
|
||||
let s = decode_bytes_to_string(&bytes).expect("UTF-16 BE BOM 应解码成功");
|
||||
assert_eq!(s, text, "UTF-16 BE BOM 文件应解码为原文本(含中文码点)");
|
||||
}
|
||||
|
||||
/// decode_bytes_to_string: 无 BOM 真 binary(密集 \x00 非 BOM 前缀) → InvalidData
|
||||
#[test]
|
||||
fn test_decode_bytes_real_binary_rejected() {
|
||||
// 真二进制:图片/编译产物常见,前两字节非 BOM 模式且含密集 NUL
|
||||
let bin: Vec<u8> = vec![0x89, 0x50, 0x4E, 0x47, 0x00, 0x0D, 0x00, 0x0A]; // PNG 头片段
|
||||
let err = decode_bytes_to_string(&bin).expect_err("真二进制应解码失败");
|
||||
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData, "真二进制返 InvalidData");
|
||||
}
|
||||
|
||||
/// decode_bytes_to_string: 普通 UTF-8 无 BOM → 透传(回归保护)
|
||||
#[test]
|
||||
fn test_decode_bytes_plain_utf8_passthrough() {
|
||||
let bytes = "普通 UTF-8 文本\n".as_bytes();
|
||||
let s = decode_bytes_to_string(bytes).expect("UTF-8 无 BOM 应直接通过");
|
||||
assert_eq!(s, "普通 UTF-8 文本\n");
|
||||
}
|
||||
|
||||
/// read_file 端到端:UTF-16 LE BOM 文件经 registry.execute 读出正常 content + binary=false
|
||||
#[tokio::test]
|
||||
async fn test_read_file_utf16_le_bom_decodes() {
|
||||
let tmp = std::env::temp_dir().join(format!("df_readfile_utf16le_{}", std::process::id()));
|
||||
let _ = fs::remove_dir_all(&tmp);
|
||||
fs::create_dir_all(&tmp).unwrap();
|
||||
// 造 UTF-16 LE BOM 文件(模拟 PowerShell Out-File 产物)
|
||||
let text = "line1\nline2\nPowerShell output\n";
|
||||
let mut bytes: Vec<u8> = vec![0xFF, 0xFE];
|
||||
for u in text.encode_utf16() {
|
||||
bytes.extend_from_slice(&u.to_le_bytes());
|
||||
}
|
||||
let file = tmp.join("ps_out.txt");
|
||||
fs::write(&file, &bytes).unwrap();
|
||||
|
||||
let mut persistent = std::collections::HashSet::new();
|
||||
persistent.insert(tmp.clone());
|
||||
let allowed_dirs = Arc::new(RwLock::new(AllowedDirs { persistent, session: Default::default(), once: Default::default() }));
|
||||
let db = Arc::new(Database::open_in_memory().await.expect("in-memory db 初始化失败"));
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""), make_test_get_app_config_ctx(&db));
|
||||
|
||||
let canon_file = file.canonicalize().unwrap().to_string_lossy().to_string();
|
||||
let args = serde_json::json!({ "path": canon_file });
|
||||
let res = registry.execute("read_file", args).await.expect("read_file 执行失败");
|
||||
// 关键断言:不再误判二进制,返正常文本内容
|
||||
assert_eq!(res["binary"], serde_json::Value::Null, "UTF-16 LE BOM 不应判二进制");
|
||||
let content = res["content"].as_str().expect("应返回 content 字段");
|
||||
assert!(content.contains("line1"), "解码内容应含 line1,实际: {}", content);
|
||||
assert!(content.contains("PowerShell output"), "解码内容应含原文,实际: {}", content);
|
||||
|
||||
fs::remove_dir_all(&tmp).ok();
|
||||
}
|
||||
|
||||
/// read_file 端到端:真 binary(密集 \x00 非 BOM)仍返 binary:true(拦截不破)
|
||||
#[tokio::test]
|
||||
async fn test_read_file_real_binary_still_rejected() {
|
||||
let tmp = std::env::temp_dir().join(format!("df_readfile_bin_{}", std::process::id()));
|
||||
let _ = fs::remove_dir_all(&tmp);
|
||||
fs::create_dir_all(&tmp).unwrap();
|
||||
// PNG 签名 + 密集 NUL(真二进制,前两字节 89 50 非 BOM)
|
||||
let bin: Vec<u8> = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0xFF, 0xFE];
|
||||
let file = tmp.join("image.png");
|
||||
fs::write(&file, &bin).unwrap();
|
||||
|
||||
let mut persistent = std::collections::HashSet::new();
|
||||
persistent.insert(tmp.clone());
|
||||
let allowed_dirs = Arc::new(RwLock::new(AllowedDirs { persistent, session: Default::default(), once: Default::default() }));
|
||||
let db = Arc::new(Database::open_in_memory().await.expect("in-memory db 初始化失败"));
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""), make_test_get_app_config_ctx(&db));
|
||||
|
||||
let canon_file = file.canonicalize().unwrap().to_string_lossy().to_string();
|
||||
let args = serde_json::json!({ "path": canon_file });
|
||||
let res = registry.execute("read_file", args).await.expect("read_file 执行失败");
|
||||
assert_eq!(res["binary"], true, "真二进制(非 BOM 前缀)应仍被拦截返 binary:true");
|
||||
assert_eq!(res["content"], serde_json::Value::Null, "真二进制 content 应为 null");
|
||||
|
||||
fs::remove_dir_all(&tmp).ok();
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// patch_file 三模式共用底层测试（F-260617-01）
|
||||
// 纯函数 apply_line_range / resolve_anchor_to_lines，不依赖 fs/async/锁。
|
||||
@@ -1921,7 +2197,7 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
let allowed_dirs = Arc::new(RwLock::new(AllowedDirs::default_with_root()));
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""));
|
||||
let registry = build_ai_tool_registry(&db, &allowed_dirs, PathBuf::from(""), make_test_get_app_config_ctx(&db));
|
||||
(db, registry)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
//! download_file AI 工具声明式注册(跨平台 URL → 文件流式下载)。
|
||||
//!
|
||||
//! 工具职责:GET 远程 URL → 流式写本地文件(reqwest bytes_stream + tokio::fs)。
|
||||
//! 替代 run_command 的 curl/wget/Invoke-WebRequest(踩 PowerShell alias/编码/跨平台陷阱)。
|
||||
//!
|
||||
//! 风险:High(下载文件副作用:写磁盘 + 触发外发网络 + 来源不可信)。SSRF 防护复用
|
||||
//! commands/ai/http.rs(validate_url / resolve_and_check_host / build_client /
|
||||
//! execute_with_redirects,经 pub(crate) 暴露)。路径校验复用 tool_registry 的
|
||||
//! validate_path + resolve_workspace_path_with_allowed(与 write_file 同款,授权目录白名单)。
|
||||
//!
|
||||
//! handler 在 commands/ai/download_file.rs::execute_download_file,声明式注册收敛样板。
|
||||
//! 闭包捕获 allowed_dirs:Arc<RwLock<AllowedDirs>>(download_file 写文件需路径授权,与 write_file 同源),
|
||||
//! 在闭包内解析 output_path(白名单快照 + symlink 逃逸拦截),把已解析的绝对路径字符串传给 handler。
|
||||
//!
|
||||
//! ## 注册位置说明
|
||||
//!
|
||||
//! download_file 既涉网络(SSRF)又涉文件(路径校验落盘),按「副作用类型」归类:
|
||||
//! 写文件(授权目录白名单)是核心安全约束,与 write_file/patch_file 同源(均持 allowed_dirs),
|
||||
//! 故在 register_file_tools 注册(tool_registry.rs),而非 register_http_tools(后者不传 allowed_dirs)。
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use df_ai::ai_tools::{AiToolRegistry, RiskLevel};
|
||||
use df_ai::declare_tool;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::commands::ai::tool_registry::resolve_workspace_path_with_allowed;
|
||||
use crate::state::AllowedDirs;
|
||||
|
||||
/// 注册 download_file 工具到 `$registry`。
|
||||
///
|
||||
/// - name/desc/schema 面向 LLM 的工具说明
|
||||
/// - risk=High(写磁盘 + 外发网络 + 来源不可信)
|
||||
/// - handler 捕获 allowed_dirs,解析 output_path 后转调 download_file.rs::execute_download_file
|
||||
/// (SSRF 防护 + 流式写 + 大小上限 + 原子写全在那)
|
||||
pub fn register(
|
||||
registry: &mut AiToolRegistry,
|
||||
allowed_dirs: &Arc<RwLock<AllowedDirs>>,
|
||||
) {
|
||||
// schema:url(必填)+ output_path(必填,落盘路径)+ timeout_secs(可选)。
|
||||
let schema = {
|
||||
let mut props = serde_json::Map::new();
|
||||
props.insert("url".into(), serde_json::json!({ "type": "string", "description": "要下载的文件 URL,仅 http/https(拒私网/localhost,SSRF 防护含 DNS resolve 后校验防重绑定,重定向≤3 跳每跳重校验)" }));
|
||||
props.insert("output_path".into(), serde_json::json!({ "type": "string", "description": "落盘文件路径(相对路径锚定到授权项目根,绝对路径须在授权目录内)。父目录不存在自动创建,覆盖既有同名文件。拒路径遍历(..)与敏感系统目录" }));
|
||||
props.insert("timeout_secs".into(), serde_json::json!({ "type": "integer", "description": "超时秒数(默认 60,上限 600)", "minimum": 1, "maximum": 600 }));
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": props,
|
||||
"required": ["url", "output_path"],
|
||||
})
|
||||
};
|
||||
declare_tool!(
|
||||
registry,
|
||||
allowed_dirs: Arc<RwLock<AllowedDirs>>,
|
||||
"download_file",
|
||||
"下载远程 URL 文件到本地磁盘(跨平台流式,替代 curl/wget/Invoke-WebRequest 的 alias/编码/跨平台陷阱)。参数:url(http/https,output_path(落盘路径,须在授权目录内),timeout_secs(默认 60,上限 600)。安全:仅 http/https + 拒私网/保留 IP(SSRF 防护)+ 路径白名单(拒遍历/越界)+ 200MB 大小上限(超限中止删半成品)+ 原子写(tmp→rename)。返回 {path, bytes_written, status, url, elapsed_ms}。下载二进制/大文件优先用此工具(不进 prompt)",
|
||||
RiskLevel::High,
|
||||
schema: schema,
|
||||
args => {
|
||||
// 路径解析:取 allowed_dirs 快照,解析 output_path(白名单 + symlink 逃逸拦截)。
|
||||
// 与 write_file 同款:相对路径锚定首个持久授权目录,绝对路径走白名单校验。
|
||||
let snap = allowed_dirs.read().await.clone();
|
||||
let output_path_raw = args["output_path"].as_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("缺少 output_path 参数"))?;
|
||||
let resolved = resolve_workspace_path_with_allowed(output_path_raw, &snap)?;
|
||||
let resolved_str = resolved.to_str()
|
||||
.ok_or_else(|| anyhow::anyhow!("output_path 含非法字符"))?
|
||||
.to_string();
|
||||
// 转调 download_file.rs handler(SSRF 防护 + 流式写 + 大小上限 + 原子写全在那)
|
||||
crate::commands::ai::download_file::execute_download_file(args, resolved_str).await
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
//! fetch_search AI 工具声明式注册(搜索引擎查询,DuckDuckGo HTML 免 key)。
|
||||
//!
|
||||
//! 工具职责:GET DuckDuckGo HTML 端点 → 解析结果列表(标题+URL+摘要)→ 截断 max_results
|
||||
//! → 返回 {query, results, count}。补「fetch_url 只抓已知 URL,无搜索能力」缺口:
|
||||
//! 用户给"调研 X"时,LLM 需先搜索找文档入口,再 fetch_url 深读。
|
||||
//!
|
||||
//! 风险:Low(只读 GET 搜索,无副作用,与 fetch_url 同级)。SSRF 防护复用 commands/ai/http.rs
|
||||
//! (validate_url / resolve_and_check_host / build_client / execute_with_redirects)。
|
||||
//! DDG HTML 免 API key(html.duckduckgo.com/html/),返回结果列表 HTML,handler 解析提取。
|
||||
//!
|
||||
//! handler 在 commands/ai/fetch_search.rs::execute_fetch_search,声明式注册收敛样板。
|
||||
//! 无 db 捕获,纯网络 GET + HTML 字符串扫描解析(对齐 fetch_url::extract_title 风格,
|
||||
//! 不引 scraper 重依赖 — html5ever+selectors+cssparser 编译产物大,DDG 单一来源不值)。
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use df_ai::ai_tools::{AiToolRegistry, RiskLevel};
|
||||
use df_ai::declare_tool;
|
||||
|
||||
/// 注册 fetch_search 工具到 `$registry`(无 db 捕获,纯网络 GET + HTML 解析)。
|
||||
///
|
||||
/// - name/desc/schema 面向 LLM 的工具说明
|
||||
/// - risk=Low(只读 GET 搜索,无副作用,与 fetch_url 同级)
|
||||
/// - handler 转调 fetch_search.rs::execute_fetch_search(SSRF 防护 + DDG GET + HTML 解析全在那)
|
||||
pub fn register(registry: &mut AiToolRegistry) {
|
||||
// 无捕获:占位 Arc<()>(handler 不持 db,仅转调 fetch_search.rs)。
|
||||
let dummy: Arc<()> = Arc::new(());
|
||||
// schema:query(必填)+ max_results(可选,默认 5,clamp [1,10])。
|
||||
let schema = {
|
||||
let mut props = serde_json::Map::new();
|
||||
props.insert("query".into(), serde_json::json!({ "type": "string", "description": "搜索关键词(必填)。返回 DuckDuckGo HTML 结果列表(标题+URL+摘要),供挑选入口用 fetch_url 深读" }));
|
||||
props.insert("max_results".into(), serde_json::json!({ "type": "integer", "description": "返回结果数(默认 5,clamp [1, 10])", "minimum": 1, "maximum": 10, "default": 5 }));
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": props,
|
||||
"required": ["query"],
|
||||
})
|
||||
};
|
||||
declare_tool!(
|
||||
registry,
|
||||
dummy: Arc<()>,
|
||||
"fetch_search",
|
||||
"搜索引擎查询(免 API key),返回 DuckDuckGo 结果列表(标题+URL+摘要)。补「fetch_url 只抓已知 URL」缺口:用户给「调研 X」时,先用本工具找文档入口,再 fetch_url 深读。参数:query(必填,搜索关键词)、max_results(默认 5,clamp [1,10])。安全:只读 GET,SSRF 防护复用 fetch_url/http_request 同套(仅 http/https + 拒私网 + 重定向≤3 跳)。返回 {query, results:[{title,url,snippet}], count, status, elapsed_ms}。每条 url 已解包 DDG 重定向跳板(uddg 参数 percent-decode),可直接 fetch_url",
|
||||
RiskLevel::Low,
|
||||
schema: schema,
|
||||
args => {
|
||||
// 转调 fetch_search.rs handler(SSRF 防护 + DDG GET + HTML 解析 + 截断全在那)
|
||||
crate::commands::ai::fetch_search::execute_fetch_search(args).await
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
//! generate_image AI 工具声明式注册(调 provider 图像生成端点 + 下载落地)。
|
||||
//!
|
||||
//! 工具职责:调 provider 的 OpenAI 兼容 `/v1/images/generations` 端点生成图片,
|
||||
//! 自动下载到工作区 output_path。默认 model=sensenova-u1-fast(SenseNova U1 信息图生成),
|
||||
//! 也支持 DALL-E 等 OpenAI 兼容图像端点。
|
||||
//!
|
||||
//! 风险:High(付费 API 按张计费 + 写磁盘 + 外发网络)。
|
||||
//!
|
||||
//! 与 download_file 的分工:download_file 是「已知 URL → 落盘」,generate_image 是
|
||||
//! 「调 provider 生成 + 自动落盘」(封装 provider 凭证 + 端点拼接 + 响应解析 + 图片下载全链路)。
|
||||
//! 图像生成模型不是 chat 模型,走 /v1/images/generations 非 chat completions,
|
||||
//! 故不能经主对话链路调用,必须用本工具显式触发。
|
||||
//!
|
||||
//! handler 在 commands/ai/generate_image.rs::execute_generate_image,声明式注册收敛样板。
|
||||
//! 闭包捕获 db:Arc<Database>(选 provider + 解析凭证)+ allowed_dirs:Arc<RwLock<AllowedDirs>>
|
||||
//! (output_path 白名单校验 + 默认路径锚定)。
|
||||
//!
|
||||
//! ## 捕获变量说明(declare_tool! 单捕获限制)
|
||||
//!
|
||||
//! declare_tool! 宏只支持单捕获变量(见 crates/df-ai/src/ai_tools_decl.rs),本工具需 db +
|
||||
//! allowed_dirs 两个,故包进 tuple Arc `(Arc<Database>, Arc<RwLock<AllowedDirs>>)` 在闭包内解构。
|
||||
//! 这是宏单捕获限制的标准绕法(零架构改动:不改宏、不改 process_tool_calls 签名)。
|
||||
//!
|
||||
//! ## 注册位置说明
|
||||
//!
|
||||
//! generate_image 涉付费 API(网络)+ 写文件(落盘 output_path),与 download_file 同源
|
||||
//! (均持 allowed_dirs),归 register_http_tools 同层(register_generate_image_tool 子函数,
|
||||
//! 在 register_http_tools 之后调用)。
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use df_ai::ai_tools::{AiToolRegistry, RiskLevel};
|
||||
use df_ai::declare_tool;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::state::AllowedDirs;
|
||||
use df_storage::db::Database;
|
||||
|
||||
/// generate_image 捕获的 (db, allowed_dirs) tuple 类型(declare_tool! 单捕获绕法)。
|
||||
type GenerateImageCtx = (Arc<Database>, Arc<RwLock<AllowedDirs>>);
|
||||
|
||||
/// 注册 generate_image 工具到 `$registry`。
|
||||
///
|
||||
/// - name/desc/schema 面向 LLM 的工具说明
|
||||
/// - risk=High(付费 API + 写磁盘 + 外发网络)
|
||||
/// - handler 捕获 (db, allowed_dirs) tuple Arc,解构后转调 generate_image.rs::execute_generate_image
|
||||
/// (provider 选择 + 凭证解析 + 端点拼接 + SSRF 防护下载 + 原子写全在那)
|
||||
pub fn register(
|
||||
registry: &mut AiToolRegistry,
|
||||
db: &Arc<Database>,
|
||||
allowed_dirs: &Arc<RwLock<AllowedDirs>>,
|
||||
) {
|
||||
// 捕获 tuple:declare_tool! 单捕获限制的绕法。
|
||||
let ctx: GenerateImageCtx = (db.clone(), allowed_dirs.clone());
|
||||
|
||||
// schema:prompt(必填)+ model(默认 sensenova-u1-fast)+ size(可选)+ n(默认 1)
|
||||
// + provider_id(可选)+ output_path(可选)。
|
||||
let schema = {
|
||||
let mut props = serde_json::Map::new();
|
||||
props.insert("prompt".into(), serde_json::json!({ "type": "string", "description": "图片描述(中英文均可),越具体生成质量越高。例:'一只戴墨镜的橘猫坐在键盘上,赛博朋克风格,霓虹光'" }));
|
||||
props.insert("model".into(), serde_json::json!({ "type": "string", "description": "图像生成模型名。默认 sensenova-u1-fast(SenseNova U1 信息图生成,走 /v1/images/generations)。也支持 dall-e-3 等 OpenAI 兼容图像模型。不传用默认", "default": "sensenova-u1-fast" }));
|
||||
props.insert("size".into(), serde_json::json!({ "type": "string", "description": "可选图片尺寸字符串,不传则让 provider 用默认。SenseNova U1 常用 2752x1536 / 1536x2752(横/竖海报),DALL-E 常用 1024x1024" }));
|
||||
props.insert("n".into(), serde_json::json!({ "type": "integer", "description": "生成数量(默认 1,clamp [1,4])。多张时取首张 url/b64 落地", "minimum": 1, "maximum": 4, "default": 1 }));
|
||||
props.insert("provider_id".into(), serde_json::json!({ "type": "string", "description": "可选,指定 provider(多 provider 时);不传则按 model 名自动匹配厂商 host(sensenova→SenseNova,dall-e→OpenAI,google/imagen→Google),匹配不到取首个 openai_compat provider" }));
|
||||
props.insert("output_path".into(), serde_json::json!({ "type": "string", "description": "可选落盘路径(相对路径锚定到授权项目根,绝对路径须在授权目录内)。不传则默认 {首个持久授权目录}/generated_images/{model}-{时间戳}.png。父目录不存在自动创建,覆盖既有同名文件。注意:这是图像生成模型(非对话模型),通过本工具调用而非主对话链路" }));
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": props,
|
||||
"required": ["prompt"],
|
||||
})
|
||||
};
|
||||
declare_tool!(
|
||||
registry,
|
||||
ctx: GenerateImageCtx,
|
||||
"generate_image",
|
||||
"调用图像生成模型(默认 SenseNova U1 Fast 信息图生成,也支持 DALL-E 等 OpenAI 兼容图像端点)生成图片并下载到工作区。参数:prompt(图片描述,必填),model(默认 sensenova-u1-fast),size(可选,SenseNova U1 常用 2752x1536/1536x2752,DALL-E 常用 1024x1024),n(默认 1,clamp [1,4]),provider_id(可选,多 provider 时指定,否则按 model 名自动匹配厂商),output_path(可选,默认 generated_images/{model}-{时间戳}.png)。注意:本工具调用图像生成模型(非对话模型),走 /v1/images/generations 端点。返回 {path, url(或 null), model, provider_id, bytes_written, elapsed_ms}。风险:付费 API 按张计费 + 写磁盘",
|
||||
RiskLevel::High,
|
||||
schema: schema,
|
||||
args => {
|
||||
// 解构 tuple 捕获:db 选 provider + 解析凭证,allowed_dirs 解析 output_path 白名单
|
||||
let (db, allowed_dirs) = ctx;
|
||||
crate::commands::ai::generate_image::execute_generate_image(args, &db, &allowed_dirs).await
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
//! get_app_config AI 工具声明式注册(查 DevFlow 自身当前 AI 配置,只读 + 脱敏)。
|
||||
//!
|
||||
//! 工具职责:返当前生效的 provider/default_provider/agent 配置/app_settings 概要 JSON,
|
||||
//! 让 LLM 一调即得「现在用什么 provider/model/迭代上限」,根治 run_command PowerShell
|
||||
//! 内联 node/python 脚本查 db(引号三重嵌套必失败)的绕行。
|
||||
//!
|
||||
//! 风险:Low(纯只读 list/load,无写副作用)。安全:api_key 经 mask_api_key 脱敏
|
||||
//! (前 4 + `••••` + 后 4),绝不返明文;app_settings 走白名单不返敏感 KV。
|
||||
//!
|
||||
//! handler body 在 commands/ai/get_app_config.rs::execute_get_app_config,
|
||||
//! 声明式注册收敛样板(对齐 fetch_url/generate_image 模式)。
|
||||
//!
|
||||
//! ## 捕获变量说明
|
||||
//!
|
||||
//! declare_tool! 宏单捕获限制:本工具需 db + 3 个 agent 配置 Arc + LlmConcurrency 共 5 项,
|
||||
//! 包进 [`GetAppConfigCtx`] struct(各字段 Arc/Clone-廉价),闭包内不解构直接
|
||||
//! `&ctx` 转调 handler。ctx 字段与 AppState 同名字段共享同一底层原子量(热改即反映)。
|
||||
|
||||
use df_ai::ai_tools::{AiToolRegistry, RiskLevel};
|
||||
use df_ai::declare_tool;
|
||||
|
||||
use crate::commands::ai::get_app_config::GetAppConfigCtx;
|
||||
|
||||
/// 注册 get_app_config 工具到 `$registry`。
|
||||
///
|
||||
/// - name/desc/schema 面向 LLM 的工具说明(无参数,返全部当前配置)
|
||||
/// - risk=Low(纯只读)
|
||||
/// - handler 捕获 GetAppConfigCtx(共享 AppState 同名 Arc 字段),转调 get_app_config.rs
|
||||
pub fn register(registry: &mut AiToolRegistry, ctx: GetAppConfigCtx) {
|
||||
// schema:无参数(object_schema 空参返仅 type:object 无 required)。
|
||||
// 不加可选 section 参数(provider/agent/settings 选返):返全部配置最简,LLM 自行读所需字段;
|
||||
// 加 section 反增参数解析分支 + schema 复杂度,违背「简洁」约束。
|
||||
let schema = df_ai::ai_tools::object_schema(vec![]);
|
||||
declare_tool!(
|
||||
registry,
|
||||
ctx: GetAppConfigCtx,
|
||||
"get_app_config",
|
||||
"获取 DevFlow 自身当前生效的 AI 配置(只读,无参数)。返回 {default_provider, providers, agent, app_settings}:default_provider 含 name/provider_type/base_url/default_model/is_default/enabled/api_key_masked(api_key 脱敏,空串=未配置 key);providers 是全部 provider 概要列表(无 api_key);agent 含 max_iterations/max_retries/approval_timeout_minutes/concurrency.per_conv(当前热改生效值);app_settings 含知识库/审批超时等 AI 相关 KV(custom_prompt 只返长度不返原文)。查「现在用什么 provider/model」「agent 迭代上限多少」「key 是否配置」用本工具,不要 run_command 执行 PowerShell 内联脚本查 db(引号嵌套易失败)。",
|
||||
RiskLevel::Low,
|
||||
schema: schema,
|
||||
_args => {
|
||||
// 转调 get_app_config.rs handler(providers/settings/agent 配置读取 + 脱敏全在那)
|
||||
crate::commands::ai::get_app_config::execute_get_app_config(&ctx).await
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -2,16 +2,22 @@
|
||||
//!
|
||||
//! 迁自 `tool_registry.rs::register_http_tools`(原 1 个 http_request),改用 `declare_tool!` 宏。
|
||||
//!
|
||||
//! 迁移策略(handler 逻辑零变更):
|
||||
//! - handler body 逐字照搬原 async move 块(转调 `crate::commands::ai::http::execute_http_request`),
|
||||
//! 仅闭包包装改由 `declare_tool!` 宏生成。
|
||||
//! 迁移策略(原 handler 逻辑零变更,2026-08-01 增 output_file/max_response_chars + allowed_dirs 捕获):
|
||||
//! - handler body 转调 `crate::commands::ai::http::execute_http_request`,闭包包装由 `declare_tool!` 宏生成。
|
||||
//! output_file 参数引入后,handler 内增 `allowed_dirs.read().await.clone()` 快照传入(路径校验用,
|
||||
//! 无 output_file 时校验不触发,行为零变更)。
|
||||
//! - name/desc/schema/risk 与原手写定义逐字一致(headers 是对象 map,object_schema 仅支持扁平
|
||||
//! 标量三元组,故保留原手工拼 serde_json::Map schema 表达式)。
|
||||
//! - 无捕获:原 handler 不 clone db(handler 仅转调 http.rs),用占位 `Arc<()>` 捕获(宏要求 capture 形参)。
|
||||
//! 标量三元组,故保留原手工拼 serde_json::Map schema 表达式);新增 output_file/max_response_chars 两 prop。
|
||||
//! - 捕获:handler 不持 db(纯转调 http.rs),但需 allowed_dirs 做 output_file 路径校验,
|
||||
//! 故捕获 `Arc<RwLock<AllowedDirs>>`(宏要求 capture 形参)。
|
||||
//!
|
||||
//! 风险:GET=Medium(只读但触发外发)/ 写方法=High(须人工批准),一个工具名两种 risk 不支持(register
|
||||
//! 单一 risk_level),故按最高风险 High 注册(写方法 High 兜底;GET 也走 High 审批更保守)。
|
||||
//!
|
||||
//! 2026-08-01 `output_file` 参数:响应可落盘(全量写不截断),路径走文件路径校验
|
||||
//! (validate_path + resolve_workspace_path_with_allowed + parent 授权,与 write_file 同源),
|
||||
//! 故 register 增 `allowed_dirs` 捕获,转调 execute_http_request 时读快照传入。
|
||||
//!
|
||||
//! 等价性验证:基线测试 `test_build_ai_tool_registry_baseline_tool_count` 仍断言 48 总量 +
|
||||
//! 工具名集合稳定(防 rename / 漏注册)。
|
||||
|
||||
@@ -19,18 +25,21 @@ use std::sync::Arc;
|
||||
|
||||
use df_ai::ai_tools::{AiToolRegistry, RiskLevel};
|
||||
use df_ai::declare_tool;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// 注册 http_request 工具到 `$registry`(无 db 捕获,纯 reqwest 调用)。
|
||||
use crate::state::AllowedDirs;
|
||||
|
||||
/// 注册 http_request 工具到 `$registry`(无 db 捕获,但捕获 allowed_dirs 供 output_file 路径校验)。
|
||||
///
|
||||
/// 与原手写 register(name, desc, schema, risk, handler) 语义 1:1:
|
||||
/// - name/desc/schema 字符串与 JSON Schema 逐字照搬原定义
|
||||
/// 与原手写 register(name, desc, schema, risk, handler) 语义 1:1(2026-08-01 增 output_file /
|
||||
/// max_response_chars 两参数 + allowed_dirs 捕获):
|
||||
/// - name/desc/schema 字符串与 JSON Schema 逐字照搬原定义(新增两 prop)
|
||||
/// - risk=High(写方法兜底,GET 同走审批更保守)
|
||||
/// - handler body 与原 async move 块逐字一致(逻辑零变更)
|
||||
/// - handler body 与原 async move 块逐字一致(逻辑零变更),新增 allowed_dirs read lock clone
|
||||
/// 传入 execute_http_request(output_file 路径校验用,无 output_file 时不触发)
|
||||
///
|
||||
/// 唯一差异:闭包包装改由 `declare_tool!` 宏生成,handler body 直接写业务逻辑。
|
||||
pub fn register(registry: &mut AiToolRegistry) {
|
||||
// 无捕获:占位 Arc<()>(原 handler 不持 db,仅转调 http.rs)。
|
||||
let dummy: Arc<()> = Arc::new(());
|
||||
pub fn register(registry: &mut AiToolRegistry, allowed_dirs: &Arc<RwLock<AllowedDirs>>) {
|
||||
// schema:headers 是对象 map,object_schema 仅支持扁平标量三元组,故手工拼 serde_json::Map。
|
||||
// 提取为 let 绑定:declare_tool! 的 schema: $schema:expr 形参对花括号块表达式解析有歧义,
|
||||
// 先求值到局部变量再传入,语义等价(原手写 register 亦以此 Map 作为 schema 实参)。
|
||||
@@ -42,6 +51,8 @@ pub fn register(registry: &mut AiToolRegistry) {
|
||||
props.insert("body".into(), serde_json::json!({ "type": "string", "description": "请求体(POST/PUT/PATCH 用),原样发送,Content-Type 须在 headers 显式指定" }));
|
||||
props.insert("timeout_secs".into(), serde_json::json!({ "type": "integer", "description": "超时秒数(默认 30,上限 60)", "minimum": 1, "maximum": 60 }));
|
||||
props.insert("parse".into(), serde_json::json!({ "type": "string", "description": "响应 body 解析:json(pretty 格式化)/text(原样)/auto(按 Content-Type 自动,默认)", "enum": ["json", "text", "auto"] }));
|
||||
props.insert("output_file".into(), serde_json::json!({ "type": "string", "description": "响应 body 写入指定文件路径(而非返回 JSON body),适合大响应。路径走文件路径校验 + 授权目录,全量写不截断。提供时返回 body=null + path/bytes_written/truncated=false" }));
|
||||
props.insert("max_response_chars".into(), serde_json::json!({ "type": "integer", "description": "body 截断阈值(字节,默认 51200/50KB,clamp [1024, 10485760])。替代默认截断上限;有 output_file 时不截断", "minimum": 1024, "maximum": 10485760 }));
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": props,
|
||||
@@ -50,14 +61,16 @@ pub fn register(registry: &mut AiToolRegistry) {
|
||||
};
|
||||
declare_tool!(
|
||||
registry,
|
||||
dummy: Arc<()>,
|
||||
allowed_dirs: Arc<RwLock<AllowedDirs>>,
|
||||
"http_request",
|
||||
"发起结构化 HTTP 请求(GET/POST/PUT/PATCH/DELETE),用于查询外部 API。参数:method(默认 GET)、url(http/https)、headers(对象 map)、body(请求体字符串)、timeout_secs(默认 30,上限 60)、parse(json/text/auto,默认 auto)。安全:仅 http/https 协议,拒绝私网/保留 IP(SSRF 防护含 DNS resolve 后校验防重绑定),重定向≤3 跳且每跳重校验。响应 body 截断 50KB。GET 为只读但触发外发网络,POST/PUT/PATCH/DELETE 有副作用,统一按高风险须人工批准",
|
||||
"发起结构化 HTTP 请求(GET/POST/PUT/PATCH/DELETE),用于查询外部 API。参数:method(默认 GET)、url(http/https)、headers(对象 map)、body(请求体字符串)、timeout_secs(默认 30,上限 60)、parse(json/text/auto,默认 auto)、output_file(响应落盘路径,适合大响应,全量写不截断)、max_response_chars(body 截断阈值字节,默认 50KB,clamp [1KB,10MB])。安全:仅 http/https 协议,拒绝私网/保留 IP(SSRF 防护含 DNS resolve 后校验防重绑定),重定向≤3 跳且每跳重校验。响应 body 默认截断 50KB(可 max_response_chars 覆盖或 output_file 落盘)。GET 为只读但触发外发网络,POST/PUT/PATCH/DELETE 有副作用,统一按高风险须人工批准",
|
||||
RiskLevel::High,
|
||||
schema: schema,
|
||||
args => {
|
||||
// 转调 http.rs handler(SSRF 防护 + 重定向 + 截断全在那)
|
||||
crate::commands::ai::http::execute_http_request(args).await
|
||||
// allowed_dirs read lock clone(output_file 路径校验用;无 output_file 时校验不触发)
|
||||
let snap = allowed_dirs.read().await.clone();
|
||||
// 转调 http.rs handler(SSRF 防护 + 重定向 + 截断 + output_file 落盘全在那)
|
||||
crate::commands::ai::http::execute_http_request(args, &snap).await
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@@ -15,7 +15,11 @@ pub mod task_graph;
|
||||
pub mod git;
|
||||
pub mod http;
|
||||
pub mod fetch_url;
|
||||
pub mod fetch_search;
|
||||
pub mod generate_image;
|
||||
pub mod get_app_config;
|
||||
pub mod workflow;
|
||||
pub mod idea;
|
||||
pub mod trash;
|
||||
pub mod file;
|
||||
pub mod download_file;
|
||||
|
||||
+26
-9
@@ -233,8 +233,28 @@ impl AppState {
|
||||
// 文件工具闭包捕获后 resolve_workspace_path_with_allowed 校验动态白名单。
|
||||
// 此处用 default_with_root 占位,下方 init 尾部 reload_allowed_dirs 从 Settings KV 覆盖。
|
||||
let allowed_dirs = Arc::new(RwLock::new(AllowedDirs::default()));
|
||||
|
||||
// get_app_config(2026-08-02):agent 配置 Arc + LlmConcurrency 先建,供 build_ai_tool_registry
|
||||
// 注入 GetAppConfigCtx(查 DevFlow 自身配置工具需读这些内存原子量,绕 AppState 自引用)。
|
||||
// 这些 Arc 不依赖 db,故可提前到 build_ai_tool_registry 之前;下方 state struct 字段复用
|
||||
// 同一 Arc(共享底层原子量,IPC 热改后 get_app_config 工具读到的即当前生效值)。
|
||||
let llm_concurrency = LlmConcurrency::new(3, 2);
|
||||
let agent_max_iterations = Arc::new(AtomicUsize::new(
|
||||
crate::commands::ai::agentic::DEFAULT_MAX_AGENT_ITERATIONS,
|
||||
));
|
||||
let agent_max_retries = Arc::new(AtomicUsize::new(
|
||||
crate::commands::ai::agentic::DEFAULT_MAX_AGENT_RETRIES,
|
||||
));
|
||||
let approval_timeout_minutes = Arc::new(AtomicU64::new(15));
|
||||
let get_app_config_ctx = crate::commands::ai::get_app_config::GetAppConfigCtx {
|
||||
db: db.clone(),
|
||||
agent_max_iterations: agent_max_iterations.clone(),
|
||||
agent_max_retries: agent_max_retries.clone(),
|
||||
approval_timeout_minutes: approval_timeout_minutes.clone(),
|
||||
llm_concurrency: llm_concurrency.clone(),
|
||||
};
|
||||
let ai_tools = Arc::new(crate::commands::ai::build_ai_tool_registry(
|
||||
&db, &allowed_dirs, data_dir.clone(),
|
||||
&db, &allowed_dirs, data_dir.clone(), get_app_config_ctx,
|
||||
));
|
||||
// Input Augmentation 层(核心设计2):ResolverRegistry 启动期注册四 resolver。
|
||||
// resolver 持 Arc<Database>(非 AppState,避免循环依赖:AppState 持 Arc<ResolverRegistry>),
|
||||
@@ -272,14 +292,11 @@ impl AppState {
|
||||
knowledge_config: Arc::new(Mutex::new(KnowledgeConfig::default())),
|
||||
data_dir: data_dir.clone(),
|
||||
settings: SettingsRepo::new(&db),
|
||||
llm_concurrency: LlmConcurrency::new(3, 2),
|
||||
agent_max_iterations: Arc::new(AtomicUsize::new(
|
||||
crate::commands::ai::agentic::DEFAULT_MAX_AGENT_ITERATIONS,
|
||||
)),
|
||||
agent_max_retries: Arc::new(AtomicUsize::new(
|
||||
crate::commands::ai::agentic::DEFAULT_MAX_AGENT_RETRIES,
|
||||
)),
|
||||
approval_timeout_minutes: Arc::new(AtomicU64::new(15)),
|
||||
// 复用上方提前构造的句柄(与 get_app_config 工具共享同一底层原子量/Semaphore)。
|
||||
llm_concurrency,
|
||||
agent_max_iterations,
|
||||
agent_max_retries,
|
||||
approval_timeout_minutes,
|
||||
workflow_state_registry: Arc::new(Mutex::new(HashMap::new())),
|
||||
// F-260619-03 Phase A: 与 ai_tools registry 共享同一 Arc(构建时注入同一句柄)
|
||||
allowed_dirs: allowed_dirs.clone(),
|
||||
|
||||
@@ -174,6 +174,15 @@ impl LlmConcurrency {
|
||||
self.per_conv.lock().await.clear();
|
||||
}
|
||||
|
||||
/// 读当前 per_conv permits 生效值(对齐 set_per_conv 的写,补足读侧 getter)。
|
||||
///
|
||||
/// 用途:get_app_config AI 工具查「当前 LLM 并发配置」时返给 LLM(global 无对称 getter:
|
||||
/// Semaphore permits 封装在 Arc<Mutex<Arc<Semaphore>>> 内层无读取接口,故此处只返 per_conv,
|
||||
/// global 不返避免误导)。热改(set_per_conv)后立即反映,无需重启。
|
||||
pub fn current_per_conv_permits(&self) -> usize {
|
||||
self.per_conv_permits.load(std::sync::atomic::Ordering::SeqCst)
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Phase3 预留(批2-B): per_sub_flow 层 — 占位未接入调用点
|
||||
// ============================================================
|
||||
|
||||
Reference in New Issue
Block a user