Skip to main content

cookie_agent_tools/
webfetch.rs

1use std::time::Duration;
2
3use async_trait::async_trait;
4use cookie_agent_engine::{
5    PreparedExecutor, PreparedTool, SessionToolContext, ToolCall, ToolError, ToolExecutionContext,
6    ToolPreparationContext, ToolProvider, ToolSpec,
7};
8use cookie_agent_protocol::{
9    ApprovalResourceSource, PermissionAction, PersistedToolResult, PreparedBindingLifetime,
10};
11use schemars::JsonSchema;
12use serde::{Deserialize, Serialize};
13
14use crate::{parse_args, prepared_operation, prepared_resource, safe_title, schema, tool_error};
15
16const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
17const DOWNLOAD_CAP: usize = 16 * 1024 * 1024;
18
19#[derive(Debug, Default)]
20pub struct WebfetchTool;
21
22#[derive(Debug, Deserialize, Serialize, JsonSchema)]
23#[serde(deny_unknown_fields)]
24struct WebfetchArgs {
25    url: String,
26    #[serde(default)]
27    raw: bool,
28}
29
30impl WebfetchArgs {
31    fn validate(&self) -> Result<(), ToolError> {
32        if !self.url.starts_with("http://") && !self.url.starts_with("https://") {
33            return Err(ToolError::InvalidUrl(
34                "URL must start with http:// or https://".into(),
35            ));
36        }
37        Ok(())
38    }
39}
40
41#[async_trait]
42impl ToolProvider for WebfetchTool {
43    fn provider_id(&self) -> &'static str {
44        "builtin.webfetch"
45    }
46
47    fn tools_for_session(&self, _ctx: &SessionToolContext) -> Result<Vec<ToolSpec>, ToolError> {
48        Ok(vec![ToolSpec {
49output: Default::default(),
50            name: "webfetch".into(),
51            permission_name: "webfetch".into(),
52            description: "Fetch an HTTP or HTTPS URL. Output has final_url, status_code, content_type, and truncated header lines, a blank line, then the text body. Returns HTML as plaintext unless raw is true; other text passes through. Use read with the artifact URI in a truncation hint to page retained output, including the body.".into(),
53            parameters: schema::<WebfetchArgs>(),
54            concurrency: cookie_agent_engine::ToolConcurrency::Parallel,
55            result_truncation: Default::default(),
56        }])
57    }
58
59    fn get_permission_name(name: &str) -> Result<&'static str, ToolError> {
60        if name != "webfetch" {
61            return Err(tool_error("webfetch provider received another tool"));
62        }
63        Ok("webfetch")
64    }
65
66    fn get_permission_resource(
67        &self,
68        name: &str,
69        arguments: &serde_json::Value,
70    ) -> Result<(&'static str, Option<String>), ToolError> {
71        let permission = Self::get_permission_name(name)?;
72        let args: WebfetchArgs = parse_args(name, arguments.clone())?;
73        args.validate()?;
74        Ok((permission, Some(args.url)))
75    }
76
77    fn get_display_argument(
78        &self,
79        name: &str,
80        arguments: &serde_json::Value,
81    ) -> Result<String, ToolError> {
82        let (_, resource) = self.get_permission_resource(name, arguments)?;
83        Ok(resource.expect("webfetch has a URL resource"))
84    }
85
86    async fn prepare(
87        &self,
88        _ctx: ToolPreparationContext,
89        call: ToolCall,
90    ) -> Result<PreparedTool, ToolError> {
91        Self::get_permission_name(&call.name)?;
92        let args: WebfetchArgs = parse_args(&call.name, call.arguments)?;
93        args.validate()?;
94        let resource = prepared_resource(
95            PermissionAction::Webfetch,
96            "url",
97            args.url.as_bytes(),
98            args.url.as_bytes(),
99            PreparedBindingLifetime::RestartStable,
100            ApprovalResourceSource::PrimaryOperation,
101        )?;
102        let operation = prepared_operation(
103            "webfetch",
104            &args,
105            vec![(PermissionAction::Webfetch, "get")],
106            vec![resource],
107            b"webfetch:v1",
108        )?;
109        let normalized = serde_json::to_value(&args).map_err(tool_error)?;
110        let label = args.url.clone();
111        PreparedTool::new(operation, normalized, None, Box::new(args))?
112            .with_policy_labels(vec![label])
113    }
114}
115
116fn request_error(error: reqwest::Error) -> ToolError {
117    if error.is_timeout() {
118        ToolError::Timeout(error.to_string())
119    } else if error.is_redirect() {
120        ToolError::RedirectError(error.to_string())
121    } else if error.is_builder() {
122        ToolError::InvalidUrl(error.to_string())
123    } else {
124        ToolError::TransportError(error.to_string())
125    }
126}
127
128impl WebfetchArgs {
129    async fn fetch(&self) -> Result<PersistedToolResult, ToolError> {
130        let client = reqwest::Client::builder()
131            .timeout(REQUEST_TIMEOUT)
132            .build()
133            .map_err(|error| ToolError::TransportError(error.to_string()))?;
134        let mut response = client.get(&self.url).send().await.map_err(request_error)?;
135        let final_url = response.url().to_string();
136        let status_code = response.status().as_u16();
137        let content_type = response
138            .headers()
139            .get(reqwest::header::CONTENT_TYPE)
140            .and_then(|value| value.to_str().ok())
141            .unwrap_or("application/octet-stream")
142            .to_owned();
143        let mime = content_type
144            .split(';')
145            .next()
146            .unwrap_or("")
147            .trim()
148            .to_ascii_lowercase();
149        let html = matches!(mime.as_str(), "text/html" | "application/xhtml+xml");
150        if !(mime.starts_with("text/")
151            || matches!(
152                mime.as_str(),
153                "application/json"
154                    | "application/xml"
155                    | "application/javascript"
156                    | "application/x-www-form-urlencoded"
157            )
158            || mime.ends_with("+json")
159            || mime.ends_with("+xml"))
160        {
161            return Err(ToolError::UnsupportedContentType(content_type));
162        }
163        let mut body = Vec::new();
164        let mut truncated = false;
165        while let Some(chunk) = response.chunk().await.map_err(request_error)? {
166            let remaining = DOWNLOAD_CAP - body.len();
167            body.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
168            if chunk.len() > remaining {
169                truncated = true;
170                break;
171            }
172        }
173        let text = if html && !self.raw {
174            html2text::from_read(body.as_slice(), 80).map_err(tool_error)?
175        } else {
176            String::from_utf8_lossy(&body).into_owned()
177        };
178        let metadata = serde_json::json!({
179            "url": self.url,
180            "final_url": final_url,
181            "status_code": status_code,
182            "content_type": content_type,
183            "truncated": truncated,
184        });
185        Ok(PersistedToolResult {
186            display: None,
187            retained_output: None,
188            title: safe_title(&self.url),
189            output: format!(
190                "final_url: {final_url}\nstatus_code: {status_code}\ncontent_type: {content_type}\ntruncated: {truncated}\n\n{text}"
191            ),
192            metadata,
193            truncation: None,
194            attachments: Vec::new(),
195            additional_messages: Vec::new(),
196        })
197    }
198}
199
200#[async_trait]
201impl PreparedExecutor for WebfetchArgs {
202    async fn revalidate(&self) -> Result<(), ToolError> {
203        Ok(())
204    }
205
206    async fn execute(
207        self: Box<Self>,
208        _context: ToolExecutionContext,
209    ) -> Result<cookie_agent_engine::ToolCompletion, ToolError> {
210        let result: Result<PersistedToolResult, ToolError> =
211            async move { self.fetch().await }.await;
212        result.map(cookie_agent_engine::ToolCompletion::single)
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use std::{
219        io::{self, BufRead, BufReader, Read, Write},
220        net::{SocketAddr, TcpListener, TcpStream},
221        sync::{
222            Arc,
223            atomic::{AtomicBool, Ordering},
224        },
225        thread::{self, JoinHandle},
226    };
227
228    use cookie_agent_protocol::{PermissionEffect, RunId, SessionId, ToolCallId};
229
230    use super::*;
231
232    struct HttpFixture {
233        address: SocketAddr,
234        stop: Arc<AtomicBool>,
235        task: Option<JoinHandle<io::Result<()>>>,
236    }
237
238    impl HttpFixture {
239        fn start() -> Self {
240            let listener = TcpListener::bind("127.0.0.1:0").unwrap();
241            let address = listener.local_addr().unwrap();
242            let stop = Arc::new(AtomicBool::new(false));
243            let task_stop = Arc::clone(&stop);
244            let task = thread::spawn(move || {
245                while !task_stop.load(Ordering::Acquire) {
246                    let (stream, _) = match listener.accept() {
247                        Ok(connection) => connection,
248                        Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
249                        Err(error) => return Err(error),
250                    };
251                    if task_stop.load(Ordering::Acquire) {
252                        break;
253                    }
254                    // Early disconnects are normal when webfetch rejects or caps a body.
255                    let _ = serve_connection(stream);
256                }
257                Ok(())
258            });
259            Self {
260                address,
261                stop,
262                task: Some(task),
263            }
264        }
265
266        fn url(&self) -> String {
267            format!("http://{}", self.address)
268        }
269    }
270
271    impl Drop for HttpFixture {
272        fn drop(&mut self) {
273            self.stop.store(true, Ordering::Release);
274            // Wake accept even if shutdown races with the loop's stop check.
275            let _ = TcpStream::connect(self.address);
276            if let Some(task) = self.task.take() {
277                task.join()
278                    .expect("join HTTP fixture")
279                    .expect("HTTP fixture accept");
280            }
281        }
282    }
283
284    fn serve_connection(mut stream: TcpStream) -> io::Result<()> {
285        stream.set_read_timeout(Some(Duration::from_secs(10)))?;
286        stream.set_write_timeout(Some(Duration::from_secs(10)))?;
287        let mut reader = BufReader::new((&mut stream).take(16 * 1024));
288        let mut request = String::new();
289        reader.read_line(&mut request)?;
290        loop {
291            let mut header = String::new();
292            if reader.read_line(&mut header)? == 0 {
293                return Err(io::ErrorKind::UnexpectedEof.into());
294            }
295            if header == "\r\n" {
296                break;
297            }
298        }
299        let Some(path) = request.split_whitespace().nth(1) else {
300            return Ok(());
301        };
302        let (status, headers, body) = match path {
303            "/first" => ("302 Found", "Location: /second\r\n", Vec::new()),
304            "/second" => ("301 Moved Permanently", "Location: /html\r\n", Vec::new()),
305            "/loop" => ("302 Found", "Location: /loop\r\n", Vec::new()),
306            "/html" => (
307                "200 OK",
308                "Content-Type: text/html; charset=utf-8\r\n",
309                b"<h1>Hello</h1><p>World <b>wide</b> web.</p>".to_vec(),
310            ),
311            "/cap" => (
312                "200 OK",
313                "Content-Type: text/plain\r\n",
314                vec![b'x'; DOWNLOAD_CAP + 1],
315            ),
316            "/exact" => (
317                "200 OK",
318                "Content-Type: text/plain\r\n",
319                vec![b'x'; DOWNLOAD_CAP],
320            ),
321            "/binary" => ("200 OK", "Content-Type: image/png\r\n", vec![0, 255]),
322            "/lines" => (
323                "200 OK",
324                "Content-Type: text/plain\r\n",
325                b"first\nsecond\r\nlast".to_vec(),
326            ),
327            _ => (
328                "404 Not Found",
329                "Content-Type: application/json\r\n",
330                b"{\"error\":\"missing\"}".to_vec(),
331            ),
332        };
333        write!(
334            stream,
335            "HTTP/1.1 {status}\r\n{headers}Content-Length: {}\r\nConnection: close\r\n\r\n",
336            body.len()
337        )?;
338        stream.write_all(&body)
339    }
340
341    async fn fetch(url: String, raw: bool) -> Result<PersistedToolResult, ToolError> {
342        let result = WebfetchArgs {
343            url: url.clone(),
344            raw,
345        }
346        .fetch()
347        .await?;
348        assert_eq!(result.metadata["url"], url);
349        let header = format!(
350            "final_url: {}\nstatus_code: {}\ncontent_type: {}\ntruncated: {}\n\n",
351            result.metadata["final_url"].as_str().unwrap(),
352            result.metadata["status_code"].as_u64().unwrap(),
353            result.metadata["content_type"].as_str().unwrap(),
354            result.metadata["truncated"].as_bool().unwrap(),
355        );
356        assert!(result.output.starts_with(&header));
357        assert_eq!(result.metadata.as_object().unwrap().len(), 5);
358        assert!(result.truncation.is_none());
359        Ok(result)
360    }
361
362    fn body(result: &PersistedToolResult) -> &str {
363        result.output.split_once("\n\n").unwrap().1
364    }
365
366    #[test]
367    fn fixture_shutdown_wakes_blocking_accept() {
368        drop(HttpFixture::start());
369    }
370
371    #[tokio::test]
372    async fn fixture_keeps_serving_after_dropped_connections() {
373        let server = HttpFixture::start();
374        for request in ["", "GET /html HTTP/1.1\r\nHost: incomplete"] {
375            let mut stream = TcpStream::connect(server.address).unwrap();
376            stream.write_all(request.as_bytes()).unwrap();
377        }
378        let result = fetch(format!("{}/lines", server.url()), false)
379            .await
380            .unwrap();
381        assert_eq!(body(&result), "first\nsecond\r\nlast");
382    }
383
384    #[test]
385    fn webfetch_keeps_standard_bounded_retention() {
386        let spec = WebfetchTool
387            .tools_for_session(&SessionToolContext::new(SessionId::new_v7()))
388            .unwrap()
389            .remove(0);
390        assert_eq!(
391            spec.result_truncation,
392            cookie_agent_engine::ToolResultTruncationPolicy::Bounded
393        );
394    }
395
396    #[tokio::test]
397    async fn redirects_preserve_initial_url_and_record_final_url() {
398        let server = HttpFixture::start();
399        let base = server.url();
400        let result = fetch(format!("{base}/first"), false).await.unwrap();
401        assert_eq!(result.metadata["url"], format!("{base}/first"));
402        assert_eq!(result.metadata["final_url"], format!("{base}/html"));
403        assert_eq!(result.metadata["status_code"], 200);
404        assert_eq!(result.metadata["content_type"], "text/html; charset=utf-8");
405        assert_eq!(result.metadata["truncated"], false);
406        let error = fetch(format!("{base}/loop"), false).await.unwrap_err();
407        assert!(matches!(error, ToolError::RedirectError(_)), "{error}");
408        let result = fetch(format!("{base}/html"), true).await.unwrap();
409        assert_eq!(result.metadata["status_code"], 200);
410    }
411
412    #[tokio::test]
413    async fn html_raw_and_plaintext_and_http_errors() {
414        let server = HttpFixture::start();
415        let base = server.url();
416        let rendered = fetch(format!("{base}/html"), false).await.unwrap();
417        let text = body(&rendered);
418        assert!(text.contains("Hello") && text.contains("World"));
419        assert!(!text.contains("<h1>") && !text.contains("<b>"));
420        let raw = fetch(format!("{base}/html"), true).await.unwrap();
421        assert_eq!(body(&raw), "<h1>Hello</h1><p>World <b>wide</b> web.</p>");
422        for raw in [false, true] {
423            let json = fetch(format!("{base}/missing?q=1"), raw).await.unwrap();
424            assert_eq!(json.metadata["status_code"], 404);
425            assert_eq!(body(&json), "{\"error\":\"missing\"}");
426            let lines = fetch(format!("{base}/lines"), raw).await.unwrap();
427            assert_eq!(body(&lines), "first\nsecond\r\nlast");
428            assert_eq!(lines.output.lines().count(), 8);
429            assert!(
430                matches!(fetch(format!("{base}/binary"), raw).await, Err(ToolError::UnsupportedContentType(content_type)) if content_type == "image/png")
431            );
432        }
433    }
434
435    #[tokio::test]
436    async fn download_cap_returns_prefix_and_only_truncates_over_cap() {
437        let server = HttpFixture::start();
438        let base = server.url();
439        for (path, truncated) in [("cap", true), ("exact", false)] {
440            let result = fetch(format!("{base}/{path}"), false).await.unwrap();
441            assert_eq!(result.metadata["truncated"], truncated);
442            assert_eq!(body(&result), "x".repeat(DOWNLOAD_CAP));
443        }
444    }
445
446    async fn prepare(arguments: serde_json::Value) -> Result<PreparedTool, ToolError> {
447        WebfetchTool
448            .prepare(
449                ToolPreparationContext {
450                    session: SessionId::new_v7(),
451                    run: RunId::new_v7(),
452                    cwd: "/tmp".into(),
453                    workspace_root: "/tmp".into(),
454                    turn_context: crate::test_turn_context(),
455                },
456                ToolCall {
457                    id: ToolCallId::new_v7(),
458                    name: "webfetch".into(),
459                    arguments,
460                },
461            )
462            .await
463    }
464
465    #[tokio::test]
466    async fn validates_scheme_and_strict_arguments_without_network() {
467        for url in [
468            "file:///tmp/a",
469            "ftp://example.org",
470            "HTTPS://example.org",
471            " example.org",
472        ] {
473            assert!(matches!(
474                prepare(serde_json::json!({"url": url})).await,
475                Err(ToolError::InvalidUrl(_))
476            ));
477        }
478        for args in [
479            serde_json::json!({}),
480            serde_json::json!({"url":"https://example.org", "raw":"true"}),
481            serde_json::json!({"url":"https://example.org", "limit":1}),
482        ] {
483            assert!(prepare(args).await.is_err());
484        }
485        let prepared = prepare(serde_json::json!({"url":"http://"})).await.unwrap();
486        assert_eq!(prepared.normalized_arguments()["raw"], false);
487        assert!(matches!(
488            fetch("http://".into(), false).await,
489            Err(ToolError::InvalidUrl(_))
490        ));
491    }
492
493    #[tokio::test]
494    async fn one_initial_url_resource_includes_query_and_defaults_to_deny() {
495        use cookie_agent_protocol::{
496            AgentDocumentSource, AgentId, AgentMode, AgentSchemaVersion, AgentSnapshot,
497            PermissionRule, Sha256Digest, WildcardPattern,
498        };
499        let server = HttpFixture::start();
500        let url = format!("{}/html?token=given", server.url());
501        let prepared = prepare(serde_json::json!({"url":url})).await.unwrap();
502        assert_eq!(prepared.operation().resources().len(), 1);
503        assert_eq!(prepared.policy_labels(), [Some(url.clone())]);
504        assert_eq!(
505            WebfetchTool
506                .get_display_argument("webfetch", prepared.normalized_arguments())
507                .unwrap(),
508            url
509        );
510        let mut policy = AgentSnapshot {
511            agent: AgentId::new("test").unwrap(),
512            schema: AgentSchemaVersion::current(),
513            mode: AgentMode::Primary,
514            description: "Test".into(),
515            document_source: AgentDocumentSource::Workspace,
516            document_fingerprint: Sha256Digest::of_bytes(b"test"),
517            composed_prompt: "test".into(),
518            prompt_fingerprint: Sha256Digest::of_bytes(b"test"),
519            max_output_tokens: 0,
520            permissions: Vec::new(),
521            delegation: None,
522            fallback_chain: Vec::new(),
523            selected_suffix_start: 0,
524        };
525        let pipeline = cookie_agent_engine::permissions::PermissionPipeline::default();
526        let decide = |policy: &AgentSnapshot| {
527            pipeline.decide_operation(
528                policy,
529                prepared.operation(),
530                prepared.policy_labels(),
531                std::path::Path::new("/tmp"),
532            )
533        };
534        assert_eq!(decide(&policy).effect, PermissionEffect::Deny);
535        policy.permissions.push(PermissionRule {
536            action: PermissionAction::Webfetch,
537            resource: WildcardPattern::new(url.split('?').next().unwrap()).unwrap(),
538            effect: PermissionEffect::Allow,
539        });
540        assert_eq!(decide(&policy).effect, PermissionEffect::Deny);
541        policy.permissions[0].resource = WildcardPattern::new(&url).unwrap();
542        for effect in [
543            PermissionEffect::Allow,
544            PermissionEffect::Ask,
545            PermissionEffect::Deny,
546        ] {
547            policy.permissions[0].effect = effect;
548            let decision = decide(&policy);
549            assert_eq!(decision.effect, effect);
550            assert_eq!(decision.evaluations.len(), 1);
551            assert_eq!(decision.evaluations[0].trace.normalized_resource, url);
552        }
553    }
554}