Skip to main content

agent_tools/special/
web_search.rs

1use crate::error::{Error, Result};
2use crate::{CmdOutput, CmdStdin, CodexOptions, CodexRequest, CodexTool, ColorMode, SandboxMode};
3use std::fs;
4use std::path::PathBuf;
5use std::time::{SystemTime, UNIX_EPOCH};
6
7pub struct WebSearchTool;
8
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct WebSearchRequest {
11    pub query: String,
12    pub timeout_ms: Option<u64>,
13    pub codex_options: CodexOptions,
14}
15
16impl WebSearchRequest {
17    pub fn new(query: impl Into<String>) -> Self {
18        let mut codex_options = CodexOptions::default();
19        codex_options.search = true;
20        codex_options.ephemeral = true;
21        codex_options.color = Some(ColorMode::Never);
22        codex_options.sandbox = Some(SandboxMode::ReadOnly);
23        codex_options.skip_git_repo_check = true;
24
25        Self {
26            query: query.into(),
27            timeout_ms: Some(120_000),
28            codex_options,
29        }
30    }
31}
32
33pub fn web_search(query: impl Into<String>) -> Result<String> {
34    let output = web_search_output(query)?;
35    Ok(output.stdout.trim().to_string())
36}
37
38pub fn web_search_output(query: impl Into<String>) -> Result<CmdOutput> {
39    WebSearchTool::search(WebSearchRequest::new(query))
40}
41
42impl WebSearchTool {
43    pub fn search(req: WebSearchRequest) -> Result<CmdOutput> {
44        let workdir = create_isolated_workdir()?;
45        let mut codex_req = CodexRequest::new(build_web_search_prompt(&req.query));
46        codex_req.stdin = Some(CmdStdin::Null);
47        codex_req.timeout_ms = req.timeout_ms;
48        codex_req.fail_on_non_zero = true;
49        codex_req.options = req.codex_options;
50        codex_req.options.cd = Some(workdir.display().to_string());
51
52        let result = CodexTool::exec(codex_req);
53        let _ = fs::remove_dir_all(workdir);
54        result
55    }
56}
57
58fn build_web_search_prompt(query: &str) -> String {
59    format!(
60        "Use live web search to answer the user's query.\n\
61         \n\
62         Requirements:\n\
63         - Search the web when needed; prefer primary, official, and recent sources.\n\
64         - If the answer depends on current information, state the date you checked.\n\
65         - Cite the main web sources used, especially for current, specific, or disputed information.\n\
66         - If sources disagree, mention the disagreement briefly.\n\
67         - Answer in the same language as the user's query unless the query asks otherwise.\n\
68         - Be concise, but include enough detail to make the answer useful.\n\
69         - Do not read, inspect, create, edit, or delete local files.\n\
70         - Do not run shell commands or use local tools unless they are strictly required for web search.\n\
71         \n\
72         User query:\n{query}"
73    )
74}
75
76fn create_isolated_workdir() -> Result<PathBuf> {
77    let unique = SystemTime::now()
78        .duration_since(UNIX_EPOCH)
79        .map(|duration| duration.as_nanos())
80        .unwrap_or_default();
81    let dir = std::env::temp_dir().join(format!(
82        "agent-tools-web-search-{}-{unique}",
83        std::process::id()
84    ));
85
86    fs::create_dir(&dir).map_err(Error::tool_io)?;
87    Ok(dir)
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn test_web_search_request_defaults() {
96        let req = WebSearchRequest::new("latest rust release");
97
98        assert_eq!(req.query, "latest rust release");
99        assert_eq!(req.timeout_ms, Some(120_000));
100        assert!(req.codex_options.search);
101        assert!(req.codex_options.ephemeral);
102        assert!(req.codex_options.skip_git_repo_check);
103        assert!(req.codex_options.model.is_none());
104        assert_eq!(req.codex_options.sandbox, Some(SandboxMode::ReadOnly));
105        assert!(req.codex_options.cd.is_none());
106    }
107
108    #[test]
109    fn test_web_search_output_uses_simple_query_defaults() {
110        let req = WebSearchRequest::new("macbook pro price");
111
112        assert_eq!(req.query, "macbook pro price");
113        assert!(req.codex_options.search);
114        assert!(req.codex_options.model.is_none());
115    }
116
117    #[test]
118    fn test_web_search_prompt_is_general() {
119        let prompt = build_web_search_prompt("最新 Rust 版本是什么?");
120
121        assert!(prompt.contains("Use live web search"));
122        assert!(prompt.contains("same language as the user's query"));
123        assert!(prompt.contains("Do not read, inspect, create, edit, or delete local files"));
124        assert!(prompt.contains("Do not run shell commands"));
125        assert!(prompt.contains("最新 Rust 版本是什么?"));
126    }
127
128    #[test]
129    fn test_isolated_workdir_is_created_outside_current_project() {
130        let dir = create_isolated_workdir().unwrap();
131
132        assert!(dir.exists());
133        assert!(dir.starts_with(std::env::temp_dir()));
134
135        fs::remove_dir_all(dir).unwrap();
136    }
137}