1use mermaid_domain::ProgressEvent;
11use std::collections::VecDeque;
12use std::sync::{Arc, Mutex, OnceLock};
13
14use async_trait::async_trait;
15use futures::{StreamExt, stream};
16
17use mermaid_domain::{FetchBackend, SearchBackend, WebConfig};
18use mermaid_domain::{ToolDefinition, ToolMetadata, ToolOutcome, ToolRunMetadata};
19
20use super::super::ctx::ExecContext;
21use super::ToolExecutor;
22use super::web_client::{
23 FetchProvider, ManagedSearxngBackend, NativeFetchClient, OllamaWebClient, SearchProvider,
24 SearxngClient, ValidatedWebUrl, WebFetchError, WebFetchResult, format_results,
25};
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum Egress {
33 OnMachine,
35 OffMachine,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq)]
44pub struct WebCapabilityStatus {
45 pub available: bool,
46 pub backend: &'static str,
47 pub trust_destination: &'static str,
48 pub egress: Egress,
49 pub reason: Option<String>,
50}
51
52impl WebCapabilityStatus {
53 #[must_use]
59 pub fn absence_reason(&self, tool: &str) -> String {
60 let reason = self.reason.as_deref().map_or_else(
61 || "backend initialization failed".to_string(),
62 mermaid_model::utils::redact_secrets,
63 );
64 let remedy = match self.backend {
65 "managed_searxng" => {
66 "The user can set [web] allow_ollama_search_fallback = true (needs \
67 OLLAMA_API_KEY), set [web] search_backend = \"ollama\", or point \
68 [web] searxng_url at a SearXNG instance, then restart mermaid"
69 },
70 "ollama_cloud" => {
71 "The user can set OLLAMA_API_KEY (or switch the [web] backend), \
72 then restart mermaid"
73 },
74 "searxng" => "The user can fix [web] searxng_url, then restart mermaid",
75 _ => "The user can check the [web] config, then restart mermaid",
76 };
77 format!(
78 "{tool} is configured to use the {} backend, which is unavailable: {reason}. {remedy}.",
79 self.backend
80 )
81 }
82}
83
84pub struct WebCapabilities {
88 pub fetch: WebCapabilityStatus,
89 pub search: WebCapabilityStatus,
90 fetch_backend: Option<Arc<dyn FetchProvider>>,
91 search_backend: Option<Arc<dyn SearchProvider>>,
92}
93
94impl WebCapabilities {
95 #[must_use]
96 pub fn resolve(web: &WebConfig) -> Self {
97 let needs_ollama_key = web.fetch_backend == FetchBackend::Ollama
98 || web.search_backend == SearchBackend::Ollama
99 || (web.search_backend == SearchBackend::Auto && web.allow_ollama_search_fallback);
100 let ollama_key = needs_ollama_key
101 .then(|| mermaid_model::utils::resolve_provider_key("ollama", "OLLAMA_API_KEY", None))
102 .flatten();
103
104 let (fetch, fetch_backend): (_, Option<Arc<dyn FetchProvider>>) = match web.fetch_backend {
105 FetchBackend::Native => match NativeFetchClient::new() {
106 Ok(client) => (
107 available("native", "direct from this machine", Egress::OnMachine),
108 Some(Arc::new(client)),
109 ),
110 Err(error) => (
111 unavailable(
112 "native",
113 "direct from this machine",
114 Egress::OnMachine,
115 error.to_string(),
116 ),
117 None,
118 ),
119 },
120 FetchBackend::Ollama => {
121 let (status, client) = ollama_cloud_backend(
122 ollama_key.clone(),
123 "Ollama Cloud (target redirects are provider-managed; final URL is not disclosed)",
124 );
125 (status, client.map(|c| c as Arc<dyn FetchProvider>))
126 },
127 };
128
129 let (search, search_backend): (_, Option<Arc<dyn SearchProvider>>) =
130 match web.search_backend {
131 SearchBackend::Auto => match crate::searxng::managed_backend_viability() {
138 Ok(_) => (
139 available(
140 "managed_searxng",
141 "local managed process",
142 Egress::OnMachine,
143 ),
144 Some(Arc::new(ManagedSearxngBackend)),
145 ),
146 Err(viability) if web.allow_ollama_search_fallback => {
147 let (status, client) = ollama_cloud_backend(ollama_key, "Ollama Cloud");
152 (
153 fallback_status(status, &viability),
154 client.map(|c| c as Arc<dyn SearchProvider>),
155 )
156 },
157 Err(reason) => (
158 unavailable(
159 "managed_searxng",
160 "local managed process",
161 Egress::OnMachine,
162 reason,
163 ),
164 None,
165 ),
166 },
167 SearchBackend::Ollama => {
168 let (status, client) = ollama_cloud_backend(ollama_key, "Ollama Cloud");
169 (status, client.map(|c| c as Arc<dyn SearchProvider>))
170 },
171 SearchBackend::Searxng => match SearxngClient::new(web.searxng_url.clone()) {
174 Ok(client) => (
175 available("searxng", "configured SearXNG instance", Egress::OffMachine),
176 Some(Arc::new(client)),
177 ),
178 Err(error) => (
179 unavailable(
180 "searxng",
181 "configured SearXNG instance",
182 Egress::OffMachine,
183 error.to_string(),
184 ),
185 None,
186 ),
187 },
188 };
189
190 Self {
191 fetch,
192 search,
193 fetch_backend,
194 search_backend,
195 }
196 }
197
198 #[cfg(test)]
202 #[must_use]
203 pub fn from_statuses_for_test(fetch: WebCapabilityStatus, search: WebCapabilityStatus) -> Self {
204 Self {
205 fetch,
206 search,
207 fetch_backend: None,
208 search_backend: None,
209 }
210 }
211
212 #[must_use]
213 pub fn fetch_tool(&self) -> Option<WebFetchTool> {
214 self.fetch_backend
215 .clone()
216 .map(|backend| WebFetchTool::new(backend, self.fetch.backend))
217 }
218
219 #[must_use]
220 pub fn search_tool(&self) -> Option<WebSearchTool> {
221 self.search_backend.clone().map(|backend| WebSearchTool {
222 backend,
223 backend_name: self.search.backend,
224 })
225 }
226}
227
228fn ollama_cloud_backend(
232 key: Option<String>,
233 trust_destination: &'static str,
234) -> (WebCapabilityStatus, Option<Arc<OllamaWebClient>>) {
235 match key {
236 Some(key) => match OllamaWebClient::new(key) {
237 Ok(client) => (
238 available("ollama_cloud", trust_destination, Egress::OffMachine),
239 Some(Arc::new(client)),
240 ),
241 Err(error) => (
242 unavailable(
243 "ollama_cloud",
244 "Ollama Cloud",
245 Egress::OffMachine,
246 error.to_string(),
247 ),
248 None,
249 ),
250 },
251 None => (
252 unavailable(
253 "ollama_cloud",
254 "Ollama Cloud",
255 Egress::OffMachine,
256 "OLLAMA_API_KEY is not configured",
257 ),
258 None,
259 ),
260 }
261}
262
263fn fallback_status(mut status: WebCapabilityStatus, viability: &str) -> WebCapabilityStatus {
269 if let Some(reason) = status.reason.take() {
270 status.reason = Some(format!(
271 "the managed bundle is unavailable ({viability}) and the configured \
272 Ollama Cloud fallback is too: {reason}"
273 ));
274 }
275 status
276}
277
278fn available(
279 backend: &'static str,
280 trust_destination: &'static str,
281 egress: Egress,
282) -> WebCapabilityStatus {
283 WebCapabilityStatus {
284 available: true,
285 backend,
286 trust_destination,
287 egress,
288 reason: None,
289 }
290}
291
292fn unavailable(
293 backend: &'static str,
294 trust_destination: &'static str,
295 egress: Egress,
296 reason: impl Into<String>,
297) -> WebCapabilityStatus {
298 WebCapabilityStatus {
299 available: false,
300 backend,
301 trust_destination,
302 egress,
303 reason: Some(reason.into()),
304 }
305}
306
307pub struct WebSearchTool {
311 backend: Arc<dyn SearchProvider>,
312 backend_name: &'static str,
313}
314
315const MAX_WEB_SEARCH_FAILURE_BYTES: usize = 1024;
316
317#[async_trait]
318impl ToolExecutor for WebSearchTool {
319 fn name(&self) -> &'static str {
320 "web_search"
321 }
322
323 fn schema(&self) -> ToolDefinition {
324 ToolDefinition {
325 name: "web_search".to_string(),
326 description:
327 "Search the web. Takes either a single `query` + `max_results`, or an array of `queries` for parallel fan-out."
328 .to_string(),
329 input_schema: serde_json::json!({
330 "type": "object",
331 "properties": {
332 "query": { "type": "string", "minLength": 1, "maxLength": 2048 },
333 "max_results": { "type": "integer", "minimum": 1, "maximum": 10, "default": 5 },
334 "queries": {
335 "type": "array",
336 "minItems": 1,
337 "maxItems": mermaid_model::constants::MAX_BATCH_TOOL_ITEMS,
338 "items": {
339 "type": "object",
340 "properties": {
341 "query": { "type": "string", "minLength": 1, "maxLength": 2048 },
342 "max_results": { "type": "integer", "minimum": 1, "maximum": 10, "default": 5 }
343 },
344 "required": ["query"],
345 "additionalProperties": false
346 }
347 }
348 },
349 "oneOf": [
350 { "required": ["query"], "not": { "required": ["queries"] } },
351 { "required": ["queries"], "not": { "required": ["query"] } }
352 ],
353 "additionalProperties": false
354 }),
355 }
356 }
357
358 #[expect(
359 clippy::too_many_lines,
360 reason = "the fan-out search: gate, run the queries concurrently, fold each result or \
361 failure into the combined text, then the all-failed error and the success both build the \
362 same ten-field WebSearch metadata from the query list and the failure list; a helper for \
363 the fold would return that pair of lists and the sources, and the outcome builders would \
364 still need everything else"
365 )]
366 async fn execute(&self, args: serde_json::Value, ctx: ExecContext) -> ToolOutcome {
367 let queries = match parse_queries(&args) {
368 Ok(q) => q,
369 Err(e) => return ToolOutcome::error(e, None),
370 };
371 if queries.is_empty() {
372 return ToolOutcome::error("web_search requires at least one query", None);
373 }
374 if let Some(blocked) = super::policy_gate::gate_external(
375 &ctx,
376 "web_search",
377 mermaid_runtime::ToolCategory::Web,
378 format!("web_search ({} queries)", queries.len()),
379 &args,
380 )
381 .await
382 {
383 return blocked;
384 }
385
386 let start = std::time::Instant::now();
387 let jobs = stream::iter(queries.iter().cloned().enumerate())
388 .map(|(idx, (query, count))| {
389 let backend = self.backend.clone();
390 let progress = ctx.progress.clone();
391 let budget = ctx.web_budget();
392 let total = queries.len();
393 async move {
394 let display_query = mermaid_model::utils::redact_secrets(&query);
395 let _ = progress
396 .send(ProgressEvent::Status(format!(
397 "searching {}/{}: {}",
398 idx + 1,
399 total,
400 display_query
401 )))
402 .await;
403 let result = backend.search(&query, count, budget).await;
404 (idx, query, result)
405 }
406 })
407 .buffer_unordered(mermaid_model::constants::MAX_WEB_SEARCH_CONCURRENCY)
408 .collect::<Vec<_>>();
409 let mut completed = tokio::select! {
410 biased;
411 _ = ctx.token.cancelled() => return ToolOutcome::cancelled(),
412 completed = jobs => completed,
413 };
414 completed.sort_by_key(|(idx, _, _)| *idx);
415
416 let mut combined = String::new();
417 let mut result_count = 0usize;
418 let mut sources = Vec::new();
419 let mut errors: Vec<mermaid_domain::WebSearchFailure> = Vec::new();
420 for (idx, query, result) in completed {
421 let display_query = mermaid_model::utils::redact_secrets(&query);
422 let section =
426 match result {
427 Ok(results) => {
428 result_count += results.len();
429 sources.extend(results.iter().map(|result| {
430 mermaid_model::utils::sanitize_url_for_display(&result.url)
431 }));
432 if results.is_empty() {
433 "[SEARCH_RESULTS]\n(no results found)\n[/SEARCH_RESULTS]\n".to_string()
434 } else {
435 format_results(&results)
436 }
437 },
438 Err(e) => {
439 let safe_error = mermaid_model::utils::truncate_middle_bytes(
440 &mermaid_model::utils::redact_secrets(&format!("{e:#}")),
441 MAX_WEB_SEARCH_FAILURE_BYTES,
442 );
443 errors.push(mermaid_domain::WebSearchFailure {
444 query_index: idx,
445 error: safe_error.clone(),
446 });
447 format!("(search failed: {safe_error})\n")
448 },
449 };
450 if queries.len() > 1 {
451 combined.push_str(&format!("=== query: {display_query} ===\n{section}\n\n"));
452 } else {
453 combined = section;
454 }
455 }
456
457 if errors.len() == queries.len() {
461 let summary = errors
462 .iter()
463 .map(|failure| format!("query {}: {}", failure.query_index + 1, failure.error))
464 .collect::<Vec<_>>()
465 .join("; ");
466 let message = format!("web_search via {} failed: {summary}", self.backend_name);
467 let message = mermaid_model::utils::truncate_middle_bytes(
468 &message,
469 mermaid_model::constants::WEB_SEARCH_AGGREGATE_MAX_BYTES
470 .saturating_sub("Error: ".len()),
471 );
472 return ToolOutcome::error(message, Some(start.elapsed().as_secs_f64())).with_metadata(
473 ToolRunMetadata {
474 detail: ToolMetadata::WebSearch {
475 queries: queries.iter().map(|(query, _)| query.clone()).collect(),
476 requested_count: queries.iter().map(|(_, count)| *count).sum(),
477 result_count: 0,
478 sources: Vec::new(),
479 backend: self.backend_name.to_string(),
480 succeeded_queries: 0,
481 failed_queries: errors.len(),
482 partial: false,
483 truncated: false,
484 failures: errors,
485 },
486 result_count: Some(0),
487 ..ToolRunMetadata::default()
488 },
489 );
490 }
491
492 let truncated = combined.len() > mermaid_model::constants::WEB_SEARCH_AGGREGATE_MAX_BYTES;
496 let combined = mermaid_model::utils::truncate_middle_bytes(
497 &combined,
498 mermaid_model::constants::WEB_SEARCH_AGGREGATE_MAX_BYTES,
499 );
500
501 let duration_secs = start.elapsed().as_secs_f64();
502 let requested_count = queries.iter().map(|(_, count)| *count).sum();
503 let query_texts = queries.iter().map(|(query, _)| query.clone()).collect();
504 ToolOutcome::success(
505 combined,
506 format!(
507 "{} {} returned",
508 result_count,
509 if result_count == 1 {
510 "result"
511 } else {
512 "results"
513 }
514 ),
515 duration_secs,
516 )
517 .with_metadata(ToolRunMetadata {
518 detail: ToolMetadata::WebSearch {
519 queries: query_texts,
520 requested_count,
521 result_count,
522 sources,
523 backend: self.backend_name.to_string(),
524 succeeded_queries: queries.len() - errors.len(),
525 failed_queries: errors.len(),
526 partial: !errors.is_empty(),
527 truncated,
528 failures: errors,
529 },
530 result_count: Some(result_count),
531 ..ToolRunMetadata::default()
532 })
533 }
534}
535
536pub struct WebFetchTool {
540 backend: Arc<dyn FetchProvider>,
541 backend_name: &'static str,
542 snapshots: Arc<Mutex<FetchSnapshotStore>>,
543}
544
545impl WebFetchTool {
546 fn new(backend: Arc<dyn FetchProvider>, backend_name: &'static str) -> Self {
547 Self {
548 backend,
549 backend_name,
550 snapshots: global_fetch_snapshot_store(),
551 }
552 }
553
554 #[cfg(test)]
555 fn new_with_test_snapshots(
556 backend: Arc<dyn FetchProvider>,
557 backend_name: &'static str,
558 ) -> Self {
559 Self {
560 backend,
561 backend_name,
562 snapshots: Arc::new(Mutex::new(FetchSnapshotStore::default())),
563 }
564 }
565}
566
567fn fetch_failure_outcome(
568 error: &WebFetchError,
569 requested_url: &str,
570 backend: &str,
571 duration_secs: f64,
572 pattern: Option<String>,
573 context_lines: usize,
574) -> ToolOutcome {
575 let requested_url = mermaid_model::utils::sanitize_url_for_display(requested_url);
576 let message = mermaid_model::utils::redact_secrets(&format!(
577 "web_fetch({requested_url}) via {backend}: {error}"
578 ));
579 let pattern_context = pattern.as_ref().map(|_| context_lines);
580 ToolOutcome::error(message, Some(duration_secs)).with_metadata(ToolRunMetadata {
581 detail: ToolMetadata::WebFetch {
582 url: requested_url,
583 final_url: None,
584 status: error.status(),
585 error_kind: Some(error.kind().to_string()),
586 media_type: None,
587 charset: None,
588 backend: backend.to_string(),
589 extraction: String::new(),
590 title: None,
591 line_count: 0,
592 byte_count: 0,
593 source_byte_count: 0,
594 output_byte_count: 0,
595 truncated: false,
596 pattern,
597 context_lines: pattern_context,
598 match_count: None,
599 snapshot_id: None,
600 },
601 line_count: Some(0),
602 byte_count: Some(0),
603 ..ToolRunMetadata::default()
604 })
605}
606
607const MAX_FETCH_SNAPSHOTS: usize = 4;
608const MAX_FETCH_SNAPSHOT_BYTES: usize = 32 * 1024 * 1024;
609const MAX_SNAPSHOT_TITLE_BYTES: usize = 300;
610const MAX_SNAPSHOT_URL_BYTES: usize = 8 * 1024;
611const MAX_SNAPSHOT_MEDIA_TYPE_BYTES: usize = 256;
612const MAX_SNAPSHOT_CHARSET_BYTES: usize = 64;
613
614static FETCH_SNAPSHOT_STORE: OnceLock<Arc<Mutex<FetchSnapshotStore>>> = OnceLock::new();
615
616#[derive(Clone, Debug, PartialEq, Eq)]
617struct FetchSnapshotScope {
618 session_id: Option<String>,
619 task_id: Option<String>,
620 fallback_turn: Option<u64>,
621}
622
623impl FetchSnapshotScope {
624 fn from_context(ctx: &ExecContext) -> Self {
625 let has_owner = ctx.session_id.is_some() || ctx.task_id.is_some();
626 Self {
627 session_id: ctx.session_id.as_deref().map(compact_string),
628 task_id: ctx.task_id.as_deref().map(compact_string),
629 fallback_turn: (!has_owner).then_some(ctx.turn.0),
632 }
633 }
634
635 fn retained_string_bytes(&self) -> usize {
636 option_string_capacity(&self.session_id)
637 .saturating_add(option_string_capacity(&self.task_id))
638 }
639}
640
641#[derive(Clone)]
642struct FetchSnapshot {
643 id: String,
644 scope: FetchSnapshotScope,
645 page: Arc<WebFetchResult>,
646 retained_bytes: usize,
647}
648
649#[derive(Default)]
650struct FetchSnapshotStore {
651 entries: VecDeque<FetchSnapshot>,
652 bytes: usize,
653 next_id: u64,
654}
655
656impl FetchSnapshotStore {
657 fn insert(
658 &mut self,
659 scope: FetchSnapshotScope,
660 page: WebFetchResult,
661 ) -> Result<(String, Arc<WebFetchResult>), String> {
662 self.next_id = self.next_id.wrapping_add(1).max(1);
663 let mut id = format!("web-{}", self.next_id);
664 id.shrink_to_fit();
665
666 let fixed_bytes = id.capacity().saturating_add(scope.retained_string_bytes());
667 if fixed_bytes >= MAX_FETCH_SNAPSHOT_BYTES {
668 return Err("web_fetch: snapshot owner identity exceeds the cache budget".to_string());
669 }
670 let page = Arc::new(bound_snapshot_page(
671 page,
672 MAX_FETCH_SNAPSHOT_BYTES - fixed_bytes,
673 ));
674 let retained_bytes = fixed_bytes.saturating_add(page_retained_string_bytes(&page));
675 if retained_bytes > MAX_FETCH_SNAPSHOT_BYTES {
676 return Err("web_fetch: snapshot metadata exceeds the cache budget".to_string());
677 }
678
679 while !self.entries.is_empty()
680 && (self.entries.len() >= MAX_FETCH_SNAPSHOTS
681 || self.bytes.saturating_add(retained_bytes) > MAX_FETCH_SNAPSHOT_BYTES)
682 {
683 if let Some(removed) = self.entries.pop_front() {
684 self.bytes = self.bytes.saturating_sub(removed.retained_bytes);
685 }
686 }
687 self.bytes = self.bytes.saturating_add(retained_bytes);
688 self.entries.push_back(FetchSnapshot {
689 id: id.clone(),
690 scope,
691 page: page.clone(),
692 retained_bytes,
693 });
694 Ok((id, page))
695 }
696
697 fn get(&self, scope: &FetchSnapshotScope, id: &str) -> Option<Arc<WebFetchResult>> {
698 self.entries
699 .iter()
700 .find(|entry| entry.id == id && &entry.scope == scope)
701 .map(|entry| Arc::clone(&entry.page))
702 }
703}
704
705fn global_fetch_snapshot_store() -> Arc<Mutex<FetchSnapshotStore>> {
706 FETCH_SNAPSHOT_STORE
707 .get_or_init(|| Arc::new(Mutex::new(FetchSnapshotStore::default())))
708 .clone()
709}
710
711fn bound_snapshot_page(mut page: WebFetchResult, max_retained_bytes: usize) -> WebFetchResult {
712 page.title = bounded_title(&page.title);
713 page.title.shrink_to_fit();
714 bound_owned_string(&mut page.requested_url, MAX_SNAPSHOT_URL_BYTES);
715 if let Some(final_url) = page.final_url.as_mut() {
716 bound_owned_string(final_url, MAX_SNAPSHOT_URL_BYTES);
717 }
718 bound_optional_string(&mut page.media_type, MAX_SNAPSHOT_MEDIA_TYPE_BYTES);
719 bound_optional_string(&mut page.charset, MAX_SNAPSHOT_CHARSET_BYTES);
720
721 let metadata_bytes = page_retained_string_bytes_without_content(&page);
722 let content_budget = max_retained_bytes.saturating_sub(metadata_bytes);
723 if page.content.len() > content_budget {
724 let cut = page.content.floor_char_boundary(content_budget);
725 page.content.truncate(cut);
726 page.truncated = true;
727 }
728 page.content.shrink_to_fit();
729 page
730}
731
732fn compact_string(value: &str) -> String {
733 let mut value = value.to_string();
734 value.shrink_to_fit();
735 value
736}
737
738fn bound_owned_string(value: &mut String, max_bytes: usize) {
739 if value.len() > max_bytes {
740 value.truncate(value.floor_char_boundary(max_bytes));
741 }
742 value.shrink_to_fit();
743}
744
745fn bound_optional_string(value: &mut Option<String>, max_bytes: usize) {
746 if let Some(value) = value {
747 bound_owned_string(value, max_bytes);
748 }
749}
750
751fn option_string_capacity(value: &Option<String>) -> usize {
752 value.as_ref().map_or(0, String::capacity)
753}
754
755fn page_retained_string_bytes_without_content(page: &WebFetchResult) -> usize {
756 page.requested_url
757 .capacity()
758 .saturating_add(option_string_capacity(&page.final_url))
759 .saturating_add(option_string_capacity(&page.media_type))
760 .saturating_add(option_string_capacity(&page.charset))
761 .saturating_add(page.title.capacity())
762}
763
764fn page_retained_string_bytes(page: &WebFetchResult) -> usize {
765 page_retained_string_bytes_without_content(page).saturating_add(page.content.capacity())
766}
767
768async fn run_snapshot_blocking<T, F>(work: F) -> Result<T, String>
769where
770 T: Send + 'static,
771 F: FnOnce() -> T + Send + 'static,
772{
773 run_snapshot_blocking_with(super::web_client::extraction_semaphore(), work).await
774}
775
776async fn run_snapshot_blocking_with<T, F>(
777 limiter: Arc<tokio::sync::Semaphore>,
778 work: F,
779) -> Result<T, String>
780where
781 T: Send + 'static,
782 F: FnOnce() -> T + Send + 'static,
783{
784 let permit = limiter
785 .acquire_owned()
786 .await
787 .map_err(|_| "web snapshot renderer is closed".to_string())?;
788 tokio::task::spawn_blocking(move || {
789 let _permit = permit;
792 work()
793 })
794 .await
795 .map_err(|error| format!("web snapshot renderer failed: {error}"))
796}
797
798#[async_trait]
799impl ToolExecutor for WebFetchTool {
800 fn name(&self) -> &'static str {
801 "web_fetch"
802 }
803
804 fn schema(&self) -> ToolDefinition {
805 ToolDefinition {
806 name: "web_fetch".to_string(),
807 description: "Fetch a public HTTP(S) URL into a bounded session snapshot, or inspect \
808 a prior snapshot without refetching. Use pattern for case-insensitive \
809 matching, or start_line + line_count for stable continuation."
810 .to_string(),
811 input_schema: serde_json::json!({
812 "type": "object",
813 "properties": {
814 "url": {
815 "type": "string",
816 "format": "uri",
817 "maxLength": 8192,
818 "description": "Public HTTP(S) URL to fetch"
819 },
820 "snapshot_id": {
821 "type": "string",
822 "pattern": "^web-[0-9]+$",
823 "description": "Snapshot returned by an earlier web_fetch call"
824 },
825 "raw": {
826 "type": "boolean",
827 "default": false,
828 "description": "With url: return the page source as served instead of \
829 the extracted main content"
830 },
831 "pattern": {
832 "type": "string",
833 "minLength": 1,
834 "maxLength": 1024,
835 "description": "Case-insensitive substring to find in the page (not a regex)"
836 },
837 "context_lines": {
838 "type": "integer",
839 "minimum": 0,
840 "maximum": 10,
841 "default": 2,
842 "description": "Context lines around each match (default 2, max 10)"
843 },
844 "start_line": {
845 "type": "integer",
846 "minimum": 1,
847 "description": "First 1-based snapshot line to return"
848 },
849 "line_count": {
850 "type": "integer",
851 "minimum": 1,
852 "maximum": 500,
853 "default": 200,
854 "description": "Maximum snapshot lines to return"
855 }
856 },
857 "oneOf": [
858 { "required": ["url"], "not": { "required": ["snapshot_id"] } },
859 { "required": ["snapshot_id"], "not": { "required": ["url"] } }
860 ],
861 "additionalProperties": false
862 }),
863 }
864 }
865
866 #[expect(
867 clippy::too_many_lines,
868 reason = "fetch-then-format: resolve a snapshot or gate and fetch a URL, then render the \
869 page and build the WebFetch metadata; both halves share the parsed request, the timer \
870 and the backend name, and the metadata literal alone is thirty lines of fields taken \
871 from the page and the format result"
872 )]
873 async fn execute(&self, args: serde_json::Value, ctx: ExecContext) -> ToolOutcome {
874 let request = match parse_fetch_args(&args) {
875 Ok(request) => request,
876 Err(error) => return ToolOutcome::error(error, None),
877 };
878 let start = std::time::Instant::now();
879 let snapshot_scope = FetchSnapshotScope::from_context(&ctx);
880 let (page, snapshot_id) = match &request.target {
881 FetchTarget::Snapshot(snapshot_id) => {
882 let page = self
883 .snapshots
884 .lock()
885 .unwrap_or_else(std::sync::PoisonError::into_inner)
886 .get(&snapshot_scope, snapshot_id);
887 let Some(page) = page else {
888 return ToolOutcome::error(
889 format!(
890 "web_fetch: snapshot '{snapshot_id}' is unavailable or was evicted"
891 ),
892 Some(start.elapsed().as_secs_f64()),
893 );
894 };
895 (page, snapshot_id.to_string())
896 },
897 FetchTarget::Url(url) => {
898 let safe_url = mermaid_model::utils::sanitize_url_for_display(url.as_str());
899 if let Some(blocked) = super::policy_gate::gate_external(
900 &ctx,
901 "web_fetch",
902 mermaid_runtime::ToolCategory::Web,
903 format!("web_fetch via {} {safe_url}", self.backend_name),
904 &args,
905 )
906 .await
907 {
908 return blocked;
909 }
910 let fetch = if request.raw {
911 self.backend.fetch_raw(url.as_str(), ctx.web_budget())
912 } else {
913 self.backend.fetch(url.as_str(), ctx.web_budget())
914 };
915 let page = tokio::select! {
916 biased;
917 _ = ctx.token.cancelled() => return ToolOutcome::cancelled(),
918 result = fetch => match result {
919 Ok(page) => page,
920 Err(error) => {
921 return fetch_failure_outcome(
922 &error,
923 url.as_str(),
924 self.backend_name,
925 start.elapsed().as_secs_f64(),
926 request.pattern.clone(),
927 request.context_lines,
928 );
929 },
930 },
931 };
932 let inserted = self
933 .snapshots
934 .lock()
935 .unwrap_or_else(std::sync::PoisonError::into_inner)
936 .insert(snapshot_scope, page);
937 match inserted {
938 Ok((snapshot_id, page)) => (page, snapshot_id),
939 Err(error) => {
940 return ToolOutcome::error(error, Some(start.elapsed().as_secs_f64()));
941 },
942 }
943 },
944 };
945
946 let render_page = Arc::clone(&page);
947 let render_snapshot_id = snapshot_id.clone();
948 let render_pattern = request.pattern.clone();
949 let render_context_lines = request.context_lines;
950 let render_start_line = request.start_line;
951 let render_line_count = request.line_count;
952 let render = run_snapshot_blocking(move || {
953 format_fetch(
954 &render_page,
955 &render_snapshot_id,
956 render_pattern.as_deref(),
957 render_context_lines,
958 render_start_line,
959 render_line_count,
960 )
961 });
962 let formatted = tokio::select! {
963 biased;
964 _ = ctx.token.cancelled() => return ToolOutcome::cancelled(),
965 result = render => match result {
966 Ok(formatted) => formatted,
967 Err(error) => {
968 return ToolOutcome::error(error, Some(start.elapsed().as_secs_f64()));
969 },
970 },
971 };
972 let duration_secs = start.elapsed().as_secs_f64();
973 let line_count = formatted.output.lines().count();
974 let byte_count = formatted.output.len();
975 let title = (!page.title.is_empty()).then(|| bounded_title(&page.title));
976 let requested_url = mermaid_model::utils::sanitize_url_for_display(&page.requested_url);
977 let final_url = page
978 .final_url
979 .as_deref()
980 .map(mermaid_model::utils::sanitize_url_for_display);
981 let pattern_context = request.pattern.as_ref().map(|_| request.context_lines);
982 ToolOutcome::success(
983 formatted.output,
984 format!(
985 "{} {} fetched via {}",
986 line_count,
987 if line_count == 1 { "line" } else { "lines" },
988 page.backend.as_str()
989 ),
990 duration_secs,
991 )
992 .with_metadata(ToolRunMetadata {
993 detail: ToolMetadata::WebFetch {
994 url: requested_url,
995 final_url,
996 status: page.status,
997 error_kind: None,
998 media_type: page.media_type.clone(),
999 charset: page.charset.clone(),
1000 backend: page.backend.as_str().to_string(),
1001 extraction: page.extraction.as_str().to_string(),
1002 title,
1003 line_count,
1004 byte_count,
1005 source_byte_count: page.source_bytes,
1006 output_byte_count: page.output_bytes,
1007 truncated: page.truncated || formatted.truncated,
1008 pattern: request.pattern,
1009 context_lines: pattern_context,
1010 match_count: formatted.match_count,
1011 snapshot_id: Some(snapshot_id),
1012 },
1013 line_count: Some(line_count),
1014 byte_count: Some(byte_count),
1015 ..ToolRunMetadata::default()
1016 })
1017 }
1018}
1019
1020enum FetchTarget {
1024 Url(ValidatedWebUrl),
1025 Snapshot(String),
1026}
1027
1028struct ParsedFetchArgs {
1029 target: FetchTarget,
1030 raw: bool,
1032 pattern: Option<String>,
1033 context_lines: usize,
1034 start_line: Option<usize>,
1035 line_count: usize,
1036}
1037
1038const FETCH_ARGS: [&str; 7] = [
1040 "url",
1041 "snapshot_id",
1042 "raw",
1043 "pattern",
1044 "context_lines",
1045 "start_line",
1046 "line_count",
1047];
1048
1049fn parse_raw(
1051 obj: &serde_json::Map<String, serde_json::Value>,
1052 target: &FetchTarget,
1053) -> Result<bool, String> {
1054 let raw = match obj.get("raw") {
1055 None => false,
1056 Some(value) => value
1057 .as_bool()
1058 .ok_or_else(|| "web_fetch: 'raw' must be a boolean".to_string())?,
1059 };
1060 if raw && matches!(target, FetchTarget::Snapshot(_)) {
1061 return Err(
1062 "web_fetch: 'raw' applies when fetching a url; a snapshot keeps the form it was \
1063 fetched in"
1064 .to_string(),
1065 );
1066 }
1067 Ok(raw)
1068}
1069
1070fn parse_fetch_args(args: &serde_json::Value) -> Result<ParsedFetchArgs, String> {
1071 let obj = args
1072 .as_object()
1073 .ok_or_else(|| "web_fetch arguments must be an object".to_string())?;
1074 for key in obj.keys() {
1075 if !FETCH_ARGS.contains(&key.as_str()) {
1076 return Err(format!("web_fetch: unknown argument '{key}'"));
1077 }
1078 }
1079
1080 let url = match obj.get("url") {
1081 None => None,
1082 Some(value) => {
1083 let raw = value
1084 .as_str()
1085 .ok_or_else(|| "web_fetch: 'url' must be a string".to_string())?
1086 .trim();
1087 if raw.len() > 8192 {
1088 return Err("web_fetch: URL exceeds 8192 bytes".to_string());
1089 }
1090 Some(ValidatedWebUrl::parse(raw).map_err(|error| format!("web_fetch: {error}"))?)
1091 },
1092 };
1093 let snapshot_id = match obj.get("snapshot_id") {
1094 None => None,
1095 Some(value) => {
1096 let id = value
1097 .as_str()
1098 .ok_or_else(|| "web_fetch: 'snapshot_id' must be a string".to_string())?;
1099 let valid = id.strip_prefix("web-").is_some_and(|suffix| {
1100 !suffix.is_empty() && suffix.bytes().all(|b| b.is_ascii_digit())
1101 });
1102 if !valid {
1103 return Err("web_fetch: invalid snapshot id".to_string());
1104 }
1105 Some(id.to_string())
1106 },
1107 };
1108 let target = match (url, snapshot_id) {
1109 (Some(url), None) => FetchTarget::Url(url),
1110 (None, Some(id)) => FetchTarget::Snapshot(id),
1111 _ => {
1112 return Err("web_fetch requires exactly one of 'url' or 'snapshot_id'".to_string());
1113 },
1114 };
1115
1116 let pattern = match obj.get("pattern") {
1117 None => None,
1118 Some(value) => {
1119 let pattern = value
1120 .as_str()
1121 .ok_or_else(|| "web_fetch: 'pattern' must be a string".to_string())?
1122 .trim();
1123 if pattern.is_empty() {
1124 return Err("web_fetch: 'pattern' must not be empty".to_string());
1125 }
1126 if pattern.contains(['\r', '\n']) {
1127 return Err("web_fetch: 'pattern' must be a single line".to_string());
1128 }
1129 if pattern.chars().count() > 1024 {
1130 return Err("web_fetch: 'pattern' exceeds 1024 characters".to_string());
1131 }
1132 Some(pattern.to_string())
1133 },
1134 };
1135 let context_lines = parse_bounded_usize(obj, "context_lines", 2, 0, 10)?;
1136 if pattern.is_none() && obj.contains_key("context_lines") {
1137 return Err("web_fetch: 'context_lines' requires 'pattern'".to_string());
1138 }
1139 let has_range = obj.contains_key("start_line") || obj.contains_key("line_count");
1140 if pattern.is_some() && has_range {
1141 return Err("web_fetch: use either 'pattern' or a line range, not both".to_string());
1142 }
1143 let start_line = has_range
1144 .then(|| parse_bounded_usize(obj, "start_line", 1, 1, usize::MAX))
1145 .transpose()?;
1146 let line_count = parse_bounded_usize(obj, "line_count", 200, 1, 500)?;
1147 let raw = parse_raw(obj, &target)?;
1148
1149 Ok(ParsedFetchArgs {
1150 target,
1151 raw,
1152 pattern,
1153 context_lines,
1154 start_line,
1155 line_count,
1156 })
1157}
1158
1159fn parse_bounded_usize(
1160 obj: &serde_json::Map<String, serde_json::Value>,
1161 key: &str,
1162 default: usize,
1163 min: usize,
1164 max: usize,
1165) -> Result<usize, String> {
1166 let Some(value) = obj.get(key) else {
1167 return Ok(default);
1168 };
1169 let value = value
1170 .as_u64()
1171 .and_then(|value| usize::try_from(value).ok())
1172 .ok_or_else(|| format!("web_fetch: '{key}' must be an integer"))?;
1173 if value < min || value > max {
1174 return Err(format!("web_fetch: '{key}' must be from {min} to {max}"));
1175 }
1176 Ok(value)
1177}
1178
1179const WEB_FETCH_MAX_BYTES: usize = mermaid_model::constants::WEB_SEARCH_AGGREGATE_MAX_BYTES;
1182const FETCH_TRUNCATION_SUFFIX: &str = "\n\n...[content truncated]\n[/WEB_FETCH]";
1183
1184struct FormattedFetch {
1185 output: String,
1186 truncated: bool,
1187 match_count: Option<usize>,
1188}
1189
1190const MAX_PATTERN_MATCHES: usize = 20;
1194
1195fn format_fetch(
1196 page: &WebFetchResult,
1197 snapshot_id: &str,
1198 pattern: Option<&str>,
1199 ctx_lines: usize,
1200 start_line: Option<usize>,
1201 line_count: usize,
1202) -> FormattedFetch {
1203 let title = if page.title.trim().is_empty() {
1204 "(no title)".to_string()
1205 } else {
1206 bounded_title(&page.title)
1207 };
1208 let requested_url = bounded_url(&page.requested_url);
1209 let final_url = page
1210 .final_url
1211 .as_deref()
1212 .map(bounded_url)
1213 .unwrap_or_else(|| "(not disclosed by backend)".to_string());
1214 let status = page
1215 .status
1216 .map(|status| status.to_string())
1217 .unwrap_or_else(|| "unknown".to_string());
1218 let media = page.media_type.as_deref().unwrap_or("unknown");
1219 let charset = page.charset.as_deref().unwrap_or("unknown");
1220
1221 let (body, match_count) = if let Some(pattern) = pattern {
1222 match extract_matches(&page.content, pattern, ctx_lines, MAX_PATTERN_MATCHES) {
1223 Some((report, count)) => (report, Some(count)),
1224 None => (format!("No matching lines for \"{pattern}\"."), Some(0)),
1225 }
1226 } else if let Some(start_line) = start_line {
1227 (
1228 format_line_range(&page.content, start_line, line_count),
1229 None,
1230 )
1231 } else {
1232 (page.content.clone(), None)
1233 };
1234
1235 let mut output = format!(
1236 "[WEB_FETCH]\nTitle: {title}\nRequested URL: {requested_url}\nFinal URL: {final_url}\nStatus: {status}\nMedia-Type: {media}\nCharset: {charset}\nBackend: {}\nExtraction: {}\nSnapshot: {snapshot_id}\nSource bytes: {}\nExtracted bytes: {}\nSnapshot bytes: {}\nSource lines: {}\n\nContent:\n{body}\n[/WEB_FETCH]",
1237 page.backend.as_str(),
1238 page.extraction.as_str(),
1239 page.source_bytes,
1240 page.output_bytes,
1241 page.content.len(),
1242 page.content.lines().count(),
1243 );
1244 let truncated = output.len() > WEB_FETCH_MAX_BYTES;
1245 if truncated {
1246 let budget = WEB_FETCH_MAX_BYTES.saturating_sub(FETCH_TRUNCATION_SUFFIX.len());
1247 let cut = output.floor_char_boundary(budget);
1248 output.truncate(cut);
1249 output.push_str(FETCH_TRUNCATION_SUFFIX);
1250 }
1251 FormattedFetch {
1252 output,
1253 truncated,
1254 match_count,
1255 }
1256}
1257
1258fn bounded_title(title: &str) -> String {
1259 let mut bounded = String::with_capacity(title.len().min(MAX_SNAPSHOT_TITLE_BYTES));
1260 let content_budget = MAX_SNAPSHOT_TITLE_BYTES.saturating_sub(3);
1261 let mut truncated = false;
1262
1263 for word in title.split_whitespace() {
1264 let separator_bytes = usize::from(!bounded.is_empty());
1265 if bounded
1266 .len()
1267 .saturating_add(separator_bytes)
1268 .saturating_add(word.len())
1269 <= MAX_SNAPSHOT_TITLE_BYTES
1270 {
1271 if separator_bytes != 0 {
1272 bounded.push(' ');
1273 }
1274 bounded.push_str(word);
1275 continue;
1276 }
1277
1278 if bounded.len() > content_budget {
1279 bounded.truncate(bounded.floor_char_boundary(content_budget));
1280 }
1281 if bounded.len() < content_budget {
1282 if separator_bytes != 0 && bounded.len() < content_budget {
1283 bounded.push(' ');
1284 }
1285 let remaining = content_budget.saturating_sub(bounded.len());
1286 let cut = word.floor_char_boundary(remaining);
1287 bounded.push_str(&word[..cut]);
1288 }
1289 truncated = true;
1290 break;
1291 }
1292
1293 if truncated {
1294 bounded.push_str("...");
1295 }
1296 bounded
1297}
1298
1299fn bounded_url(url: &str) -> String {
1300 const MAX_DISPLAY_URL_BYTES: usize = 2048;
1301 let url = mermaid_model::utils::sanitize_url_for_display(url);
1302 if url.len() <= MAX_DISPLAY_URL_BYTES {
1303 return url;
1304 }
1305 let cut = url.floor_char_boundary(MAX_DISPLAY_URL_BYTES.saturating_sub(3));
1306 format!("{}...", &url[..cut])
1307}
1308
1309fn format_line_range(content: &str, start_line: usize, line_count: usize) -> String {
1310 let total = content.lines().count();
1311 if start_line > total {
1312 return format!("Requested line {start_line}, but the snapshot contains {total} lines.");
1313 }
1314 let mut output = format!(
1315 "Lines {start_line}-{} of {total}:\n",
1316 start_line
1317 .saturating_add(line_count)
1318 .saturating_sub(1)
1319 .min(total)
1320 );
1321 for (offset, line) in content
1322 .lines()
1323 .skip(start_line.saturating_sub(1))
1324 .take(line_count)
1325 .enumerate()
1326 {
1327 output.push_str(&format!("L{}: {line}\n", start_line + offset));
1328 }
1329 output
1330}
1331
1332fn normalized_case_fold(value: &str) -> String {
1337 use caseless::Caseless;
1338 use unicode_normalization::UnicodeNormalization;
1339
1340 value.nfd().default_case_fold().nfd().collect()
1341}
1342
1343fn extract_matches(
1350 content: &str,
1351 pattern: &str,
1352 context_lines: usize,
1353 max_blocks: usize,
1354) -> Option<(String, usize)> {
1355 let needle = normalized_case_fold(pattern);
1356 let lines: Vec<&str> = content.lines().collect();
1357 let matched: Vec<usize> = lines
1358 .iter()
1359 .enumerate()
1360 .filter(|(_, line)| normalized_case_fold(line).contains(&needle))
1361 .map(|(i, _)| i)
1362 .collect();
1363 if matched.is_empty() {
1364 return None;
1365 }
1366
1367 let mut blocks: Vec<(usize, usize)> = Vec::new();
1370 for &i in &matched {
1371 let start = i.saturating_sub(context_lines);
1372 let end = (i + context_lines).min(lines.len() - 1);
1373 match blocks.last_mut() {
1374 Some((_, last_end)) if start <= *last_end + 1 => *last_end = (*last_end).max(end),
1375 _ => blocks.push((start, end)),
1376 }
1377 }
1378 let included = &blocks[..blocks.len().min(max_blocks)];
1379 let cutoff = included.last().map(|&(_, end)| end).unwrap_or(0);
1380 let dropped = matched.iter().filter(|&&i| i > cutoff).count();
1381
1382 let mut out = format!(
1383 "{} match{} for \"{}\":\n",
1384 matched.len(),
1385 if matched.len() == 1 { "" } else { "es" },
1386 pattern
1387 );
1388 for (bi, &(start, end)) in included.iter().enumerate() {
1389 if bi > 0 {
1390 out.push_str("---\n");
1391 }
1392 for (offset, line) in lines[start..=end].iter().enumerate() {
1393 out.push_str(&format!("L{}: {}\n", start + offset + 1, line));
1395 }
1396 }
1397 if dropped > 0 {
1398 out.push_str(&format!(
1399 "(+{dropped} more match{})\n",
1400 if dropped == 1 { "" } else { "es" }
1401 ));
1402 }
1403 let match_count = matched.len();
1404 Some((out, match_count))
1405}
1406
1407fn parse_queries(args: &serde_json::Value) -> Result<Vec<(String, usize)>, String> {
1408 let obj = args
1409 .as_object()
1410 .ok_or_else(|| "web_search arguments must be an object".to_string())?;
1411 for key in obj.keys() {
1412 if !matches!(key.as_str(), "query" | "max_results" | "queries") {
1413 return Err(format!("web_search: unknown argument '{key}'"));
1414 }
1415 }
1416 if obj.contains_key("query") && obj.contains_key("queries") {
1417 return Err("web_search accepts either 'query' or 'queries', not both".to_string());
1418 }
1419
1420 if let Some(value) = obj.get("queries") {
1421 let Some(arr) = value.as_array() else {
1422 return Err("web_search: 'queries' must be an array".to_string());
1423 };
1424 if arr.is_empty() {
1425 return Err("web_search: 'queries' must contain at least one entry".to_string());
1426 }
1427 if arr.len() > mermaid_model::constants::MAX_BATCH_TOOL_ITEMS {
1428 return Err(format!(
1429 "web_search: too many queries ({}); cap is {} per call — split the request",
1430 arr.len(),
1431 mermaid_model::constants::MAX_BATCH_TOOL_ITEMS
1432 ));
1433 }
1434 let mut out = Vec::with_capacity(arr.len());
1435 for v in arr {
1436 let Some(obj) = v.as_object() else {
1437 return Err(
1438 "web_search: 'queries' must be an array of {query, max_results}".to_string(),
1439 );
1440 };
1441 for key in obj.keys() {
1442 if !matches!(key.as_str(), "query" | "max_results") {
1443 return Err(format!("web_search: unknown query argument '{key}'"));
1444 }
1445 }
1446 out.push(parse_query_entry(obj)?);
1447 }
1448 return Ok(out);
1449 }
1450 if obj.contains_key("query") {
1451 return Ok(vec![parse_query_entry(obj)?]);
1452 }
1453 Err("web_search requires 'query' (string) or 'queries' (array)".to_string())
1454}
1455
1456fn parse_query_entry(
1457 obj: &serde_json::Map<String, serde_json::Value>,
1458) -> Result<(String, usize), String> {
1459 let query = obj
1460 .get("query")
1461 .and_then(|value| value.as_str())
1462 .ok_or_else(|| "web_search: each query needs 'query' (string)".to_string())?
1463 .trim();
1464 if query.is_empty() {
1465 return Err("web_search: query must not be empty".to_string());
1466 }
1467 if query.contains(['\r', '\n', '\0']) {
1468 return Err("web_search: query must be a single text line".to_string());
1469 }
1470 if query.chars().count() > 2048 {
1471 return Err("web_search: query exceeds 2048 characters".to_string());
1472 }
1473 let count = match obj.get("max_results") {
1474 None => 5,
1475 Some(value) => {
1476 let count = value.as_u64().ok_or_else(|| {
1477 "web_search: 'max_results' must be an integer from 1 to 10".to_string()
1478 })?;
1479 if !(1..=10).contains(&count) {
1480 return Err("web_search: 'max_results' must be from 1 to 10".to_string());
1481 }
1482 count as usize
1483 },
1484 };
1485 Ok((query.to_string(), count))
1486}
1487
1488pub(crate) fn require_http_scheme(url: &str) -> Result<reqwest::Url, String> {
1500 let parsed = reqwest::Url::parse(url).map_err(|e| format!("invalid URL: {e}"))?;
1501 match parsed.scheme() {
1502 "http" | "https" => Ok(parsed),
1503 other => Err(format!(
1504 "unsupported URL scheme '{other}' (only http/https allowed)"
1505 )),
1506 }
1507}
1508
1509#[cfg(test)]
1510mod tests {
1511 use super::*;
1512
1513 fn page(content: impl Into<String>) -> WebFetchResult {
1514 let content = content.into();
1515 WebFetchResult {
1516 requested_url: "https://example.com/start".to_string(),
1517 final_url: Some("https://example.com/final".to_string()),
1518 status: Some(200),
1519 media_type: Some("text/html".to_string()),
1520 charset: Some("utf-8".to_string()),
1521 backend: super::super::web_client::FetchBackend::Native,
1522 extraction: super::super::web_client::ExtractionMode::Readability,
1523 source_bytes: content.len(),
1524 output_bytes: content.len(),
1525 truncated: false,
1526 title: "T".to_string(),
1527 content,
1528 }
1529 }
1530
1531 fn scope(session_id: &str) -> FetchSnapshotScope {
1532 FetchSnapshotScope {
1533 session_id: Some(compact_string(session_id)),
1534 task_id: None,
1535 fallback_turn: None,
1536 }
1537 }
1538
1539 #[test]
1540 fn require_http_scheme_accepts_http_rejects_exotic() {
1541 for good in [
1544 "http://example.com",
1545 "https://example.com/path?a=1&b=2",
1546 "http://localhost:3000",
1547 "http://127.0.0.1:8080",
1548 ] {
1549 assert!(require_http_scheme(good).is_ok(), "{good} should pass");
1550 }
1551 for bad in [
1553 "file:///etc/passwd",
1554 "javascript:alert(1)",
1555 "data:text/html,<script>",
1556 "ftp://example.com",
1557 "not a url",
1558 ] {
1559 assert!(
1560 require_http_scheme(bad).is_err(),
1561 "{bad} should be rejected"
1562 );
1563 }
1564 }
1565
1566 #[test]
1567 fn format_fetch_caps_long_content() {
1568 let big = "z".repeat(WEB_FETCH_MAX_BYTES * 2);
1570 let big_page = page(big);
1571 let out = format_fetch(&big_page, "web-1", None, 2, None, 200);
1572 assert!(
1573 out.output.len() <= WEB_FETCH_MAX_BYTES,
1574 "content must be capped, got {} bytes",
1575 out.output.len()
1576 );
1577 assert!(
1578 out.output.contains("truncated"),
1579 "expected truncation marker"
1580 );
1581 assert!(out.truncated);
1582
1583 let small = page("hello world");
1585 let out = format_fetch(&small, "web-1", None, 2, None, 200);
1586 assert!(out.output.contains("hello world"));
1587 assert!(!out.output.contains("truncated"));
1588 }
1589
1590 #[test]
1591 fn format_fetch_caps_the_complete_envelope_and_sanitizes_provenance() {
1592 let mut page = page("body");
1593 page.title = format!(" {}\n{} ", "title ".repeat(100), "tail");
1594 page.requested_url = format!(
1595 "https://alice:hunter2@example.com/page?token=opaque-secret&q={}",
1596 "x".repeat(10_000)
1597 );
1598 page.final_url = Some(page.requested_url.clone());
1599
1600 let out = format_fetch(&page, "web-1", None, 2, None, 200);
1601 assert!(out.output.len() <= WEB_FETCH_MAX_BYTES);
1602 assert!(
1603 !out.output.contains("alice"),
1604 "userinfo leaked: {}",
1605 out.output
1606 );
1607 assert!(
1608 !out.output.contains("hunter2"),
1609 "password leaked: {}",
1610 out.output
1611 );
1612 assert!(!out.output.contains("opaque-secret"), "query secret leaked");
1613 let title = out
1614 .output
1615 .lines()
1616 .find_map(|line| line.strip_prefix("Title: "))
1617 .expect("title header");
1618 assert!(title.len() <= 300);
1619 }
1620
1621 #[test]
1622 fn complete_output_budget_holds_for_multibyte_boundary_sizes() {
1623 for unit in ["a", "é", "界"] {
1624 for units in [0, 1, 14_900, 15_000, 15_100, 40_000] {
1625 let mut candidate = page(unit.repeat(units));
1626 candidate.title = unit.repeat(1_000);
1627 let formatted = format_fetch(&candidate, "web-99", None, 2, None, 200);
1628 assert!(
1629 formatted.output.len() <= WEB_FETCH_MAX_BYTES,
1630 "{} bytes escaped the complete-result cap",
1631 formatted.output.len()
1632 );
1633 assert!(std::str::from_utf8(formatted.output.as_bytes()).is_ok());
1634 assert!(formatted.output.ends_with("[/WEB_FETCH]"));
1635 }
1636 }
1637 }
1638
1639 #[test]
1640 fn snapshot_store_accounts_for_and_bounds_every_retained_string() {
1641 let original_content = "é".repeat(MAX_FETCH_SNAPSHOT_BYTES / 2 + 1_000);
1642 let original_output_bytes = original_content.len();
1643 let mut oversized = page(original_content);
1644 oversized.output_bytes = original_output_bytes;
1645 oversized.title = "title ".repeat(10_000);
1646 oversized.requested_url = format!("https://example.com/{}", "r".repeat(20_000));
1647 oversized.final_url = Some(format!("https://example.com/{}", "f".repeat(20_000)));
1648 oversized.media_type = Some("m".repeat(1_000));
1649 oversized.charset = Some("c".repeat(1_000));
1650
1651 let owner = scope("session-retained-size");
1652 let mut store = FetchSnapshotStore::default();
1653 let (id, bounded) = store.insert(owner.clone(), oversized).unwrap();
1654 let entry = store.entries.back().expect("snapshot entry");
1655 let expected = entry
1656 .id
1657 .capacity()
1658 .saturating_add(entry.scope.retained_string_bytes())
1659 .saturating_add(page_retained_string_bytes(&entry.page));
1660
1661 assert_eq!(entry.id, id);
1662 assert_eq!(entry.retained_bytes, expected);
1663 assert_eq!(store.bytes, expected);
1664 assert!(store.bytes <= MAX_FETCH_SNAPSHOT_BYTES);
1665 assert_eq!(bounded.output_bytes, original_output_bytes);
1666 assert!(bounded.truncated);
1667 assert!(bounded.title.len() <= MAX_SNAPSHOT_TITLE_BYTES);
1668 assert!(bounded.requested_url.len() <= MAX_SNAPSHOT_URL_BYTES);
1669 assert!(bounded.final_url.as_ref().unwrap().len() <= MAX_SNAPSHOT_URL_BYTES);
1670 assert!(bounded.media_type.as_ref().unwrap().len() <= MAX_SNAPSHOT_MEDIA_TYPE_BYTES);
1671 assert!(bounded.charset.as_ref().unwrap().len() <= MAX_SNAPSHOT_CHARSET_BYTES);
1672 assert!(std::str::from_utf8(bounded.content.as_bytes()).is_ok());
1673 }
1674
1675 #[test]
1676 fn snapshot_store_isolates_session_and_task_owners() {
1677 let owner = FetchSnapshotScope {
1678 session_id: Some(compact_string("session-a")),
1679 task_id: Some(compact_string("task-a")),
1680 fallback_turn: None,
1681 };
1682 let mut store = FetchSnapshotStore::default();
1683 let (id, _) = store.insert(owner.clone(), page("private page")).unwrap();
1684
1685 assert!(store.get(&owner, &id).is_some());
1686 for outsider in [
1687 FetchSnapshotScope {
1688 session_id: Some(compact_string("session-b")),
1689 task_id: Some(compact_string("task-a")),
1690 fallback_turn: None,
1691 },
1692 FetchSnapshotScope {
1693 session_id: Some(compact_string("session-a")),
1694 task_id: Some(compact_string("task-b")),
1695 fallback_turn: None,
1696 },
1697 ] {
1698 assert!(store.get(&outsider, &id).is_none());
1699 }
1700 }
1701
1702 #[test]
1703 fn snapshot_store_evicts_the_oldest_entry_at_the_count_limit() {
1704 let mut store = FetchSnapshotStore::default();
1705 let owner = scope("session-eviction");
1706 let mut ids = Vec::new();
1707 for index in 0..=MAX_FETCH_SNAPSHOTS {
1708 ids.push(
1709 store
1710 .insert(owner.clone(), page(format!("page {index}")))
1711 .unwrap()
1712 .0,
1713 );
1714 }
1715 assert!(
1716 store.get(&owner, &ids[0]).is_none(),
1717 "oldest snapshot was not evicted"
1718 );
1719 assert!(store.get(&owner, ids.last().unwrap()).is_some());
1720 assert_eq!(store.entries.len(), MAX_FETCH_SNAPSHOTS);
1721 }
1722
1723 #[test]
1724 fn extract_matches_finds_case_insensitive_with_context() {
1725 let content = "line one\nline two\nTARGET here\nline four\nline five";
1726 let (out, count) = extract_matches(content, "target", 1, 20).unwrap();
1727 assert_eq!(count, 1);
1728 assert!(out.starts_with("1 match for \"target\":"));
1729 assert!(out.contains("L2: line two"));
1730 assert!(out.contains("L3: TARGET here"));
1731 assert!(out.contains("L4: line four"));
1732 assert!(!out.contains("L1:"), "context clipped to 1 line: {out}");
1733 assert!(!out.contains("L5:"));
1734 }
1735
1736 #[test]
1737 fn extract_matches_merges_overlapping_windows() {
1738 let content = "a\nhit one\nhit two\nb\nc\nd\ne\nf\ng\nhit three\nz";
1740 let (out, count) = extract_matches(content, "hit", 1, 20).unwrap();
1741 assert_eq!(count, 3);
1742 assert!(out.starts_with("3 matches"));
1743 assert_eq!(out.matches("---").count(), 1, "two blocks: {out}");
1744 assert_eq!(out.matches("hit one").count(), 1);
1746 }
1747
1748 #[test]
1749 fn extract_matches_caps_blocks_and_reports_tail() {
1750 let content = (0..25)
1752 .map(|i| format!("match {i}\nx\nx\nx\nx\nx"))
1753 .collect::<Vec<_>>()
1754 .join("\n");
1755 let (out, count) = extract_matches(&content, "match", 0, 20).unwrap();
1756 assert_eq!(count, 25);
1757 assert!(out.starts_with("25 matches"));
1758 assert_eq!(out.matches("---").count(), 19, "20 blocks: {out}");
1759 assert!(out.contains("(+5 more matches)"), "tail note: {out}");
1760 }
1761
1762 #[test]
1763 fn extract_matches_none_and_multibyte() {
1764 assert!(extract_matches("nothing here", "absent", 2, 20).is_none());
1765 let content = "voil\u{e0} un r\u{e9}sultat\nplain line";
1767 let (out, count) = extract_matches(content, "R\u{c9}SULTAT", 0, 20).unwrap();
1768 assert_eq!(count, 1);
1769 assert!(out.contains("L1: voil\u{e0} un r\u{e9}sultat"));
1770 assert!(!out.contains("plain line"));
1772 }
1773
1774 #[test]
1775 fn extract_matches_uses_full_unicode_case_folding() {
1776 let content = "Die Straße ist lang\nSTRASSE in capitals\nother";
1777 let (out, count) = extract_matches(content, "strasse", 0, 20).unwrap();
1778 assert!(out.starts_with("2 matches for \"strasse\":"), "{out}");
1779 assert!(out.contains("L1: Die Straße ist lang"), "{out}");
1780 assert!(out.contains("L2: STRASSE in capitals"), "{out}");
1781 assert_eq!(count, 2);
1782 }
1783
1784 #[test]
1785 fn extract_matches_normalizes_composed_and_decomposed_text() {
1786 let content = "Café noir\nCafe\u{301} blanc\nplain";
1787 let decomposed_pattern = "CAFE\u{301}";
1788 let (out, count) = extract_matches(content, decomposed_pattern, 0, 20).unwrap();
1789 assert!(out.starts_with("2 matches"), "{out}");
1790 assert!(out.contains("L1: Café noir"), "{out}");
1791 assert!(out.contains("L2: Cafe\u{301} blanc"), "{out}");
1792 assert_eq!(count, 2);
1793 }
1794
1795 #[test]
1796 fn format_fetch_pattern_paths() {
1797 let page = page("alpha\nbeta\ngamma");
1798 let out = format_fetch(&page, "web-1", Some("beta"), 1, None, 200);
1800 assert!(out.output.contains("1 match for \"beta\""));
1801 assert!(out.output.contains("L2: beta"));
1802 let out = format_fetch(&page, "web-1", Some("nope"), 1, None, 200);
1804 assert!(out.output.contains("No matching lines for \"nope\"."));
1805 assert!(!out.output.contains("alpha"));
1806 }
1807
1808 #[test]
1809 fn find_in_page_runs_before_the_cap() {
1810 let mut content = "x\n".repeat(WEB_FETCH_MAX_BYTES / 2);
1813 content.push_str("needle in the tail\n");
1814 let page = page(content);
1815 let out = format_fetch(&page, "web-1", Some("needle"), 1, None, 200);
1816 assert!(
1817 out.output.contains("1 match for \"needle\""),
1818 "tail match found"
1819 );
1820 assert!(out.output.contains("needle in the tail"));
1821 }
1822
1823 #[test]
1824 fn parse_queries_single_form() {
1825 let args = serde_json::json!({"query": "rust async", "max_results": 3});
1826 let q = parse_queries(&args).unwrap();
1827 assert_eq!(q.len(), 1);
1828 assert_eq!(q[0].0, "rust async");
1829 assert_eq!(q[0].1, 3);
1830 }
1831
1832 #[test]
1833 fn parse_queries_array_form() {
1834 let args = serde_json::json!({"queries": [
1835 {"query": "a", "max_results": 2},
1836 {"query": "b", "max_results": 5},
1837 ]});
1838 let q = parse_queries(&args).unwrap();
1839 assert_eq!(q.len(), 2);
1840 assert_eq!(q[1].1, 5);
1841 }
1842
1843 #[test]
1844 fn parse_queries_missing_errors() {
1845 let args = serde_json::json!({});
1846 assert!(parse_queries(&args).is_err());
1847 }
1848
1849 #[test]
1850 fn parse_queries_rejects_out_of_range_count() {
1851 let args = serde_json::json!({"query": "q", "max_results": 999});
1852 assert!(parse_queries(&args).is_err());
1853 let args = serde_json::json!({"query": "q", "max_results": 0});
1854 assert!(parse_queries(&args).is_err());
1855 let args = serde_json::json!({"query": "q", "max_results": "5"});
1856 assert!(parse_queries(&args).is_err());
1857 }
1858
1859 #[test]
1860 fn parse_queries_rejects_ambiguous_and_unknown_arguments() {
1861 assert!(
1862 parse_queries(&serde_json::json!({"query":"a", "queries":[{"query":"b"}]})).is_err()
1863 );
1864 assert!(parse_queries(&serde_json::json!({"query":"a", "extra":true})).is_err());
1865 assert!(parse_queries(&serde_json::json!({"query":" "})).is_err());
1866 assert!(parse_queries(&serde_json::json!({"query":"safe\n=== injected ==="})).is_err());
1867 }
1868
1869 #[test]
1870 fn parse_queries_rejects_excess_fan_out() {
1871 let many: Vec<_> = (0..mermaid_model::constants::MAX_BATCH_TOOL_ITEMS + 1)
1873 .map(|i| serde_json::json!({"query": format!("q{i}")}))
1874 .collect();
1875 let args = serde_json::json!({ "queries": many });
1876 assert!(parse_queries(&args).is_err());
1877
1878 let at_cap: Vec<_> = (0..mermaid_model::constants::MAX_BATCH_TOOL_ITEMS)
1880 .map(|i| serde_json::json!({"query": format!("q{i}")}))
1881 .collect();
1882 let args = serde_json::json!({ "queries": at_cap });
1883 assert_eq!(
1884 parse_queries(&args).unwrap().len(),
1885 mermaid_model::constants::MAX_BATCH_TOOL_ITEMS
1886 );
1887 }
1888
1889 #[test]
1890 fn parse_fetch_args_is_strict_and_rejects_credentialed_urls() {
1891 for invalid in [
1892 serde_json::json!({}),
1893 serde_json::json!({"url": "https://example.com", "snapshot_id": "web-1"}),
1894 serde_json::json!({"url": "https://user:password@example.com"}),
1895 serde_json::json!({"url": "http://127.0.0.1/private"}),
1896 serde_json::json!({"snapshot_id": "bad"}),
1897 serde_json::json!({"snapshot_id": "web-1", "context_lines": 2}),
1898 serde_json::json!({"snapshot_id": "web-1", "pattern": "x", "start_line": 1}),
1899 serde_json::json!({"snapshot_id": "web-1", "unknown": true}),
1900 serde_json::json!({"snapshot_id": "web-1", "raw": true}),
1901 serde_json::json!({"url": "https://example.com", "raw": "yes"}),
1902 ] {
1903 assert!(parse_fetch_args(&invalid).is_err(), "accepted {invalid}");
1904 }
1905
1906 let parsed = parse_fetch_args(&serde_json::json!({
1907 "url": "https://example.com/page#fragment",
1908 "start_line": 4,
1909 "line_count": 2
1910 }))
1911 .unwrap();
1912 let FetchTarget::Url(url) = &parsed.target else {
1913 panic!("expected a URL target");
1914 };
1915 assert_eq!(url.as_str(), "https://example.com/page");
1916 assert_eq!(parsed.start_line, Some(4));
1917 assert_eq!(parsed.line_count, 2);
1918 assert!(!parsed.raw);
1919 let raw = parse_fetch_args(&serde_json::json!({
1920 "url": "https://example.com/page",
1921 "raw": true
1922 }))
1923 .unwrap();
1924 assert!(raw.raw);
1925 }
1926
1927 #[tokio::test]
1930 async fn raw_fetch_needs_a_backend_that_can_serve_it() {
1931 use crate::providers::ctx::test_exec_context;
1932 use mermaid_domain::{ToolCallId, TurnId};
1933
1934 struct ExtractOnly;
1935
1936 #[async_trait]
1937 impl FetchProvider for ExtractOnly {
1938 async fn fetch(
1939 &self,
1940 url: &str,
1941 budget: crate::providers::ctx::WebByteBudget,
1942 ) -> Result<WebFetchResult, WebFetchError> {
1943 let _ = budget;
1944 let mut result = page("extracted");
1945 result.requested_url = url.to_string();
1946 Ok(result)
1947 }
1948 }
1949
1950 let tool = WebFetchTool::new_with_test_snapshots(Arc::new(ExtractOnly), "mock");
1951 let (ctx, _rx) =
1952 test_exec_context(TurnId(9), ToolCallId(9), std::path::PathBuf::from("/tmp"));
1953 let outcome = tool
1954 .execute(
1955 serde_json::json!({"url": "https://example.com/page", "raw": true}),
1956 ctx,
1957 )
1958 .await;
1959 assert!(!outcome.is_success());
1960 assert!(outcome.output().contains("native"), "{}", outcome.output());
1961 }
1962
1963 #[tokio::test]
1964 async fn web_fetch_failure_retains_typed_backend_provenance() {
1965 use crate::providers::ctx::test_exec_context;
1966 use mermaid_domain::{ToolCallId, ToolStatus, TurnId};
1967
1968 struct FailingFetch;
1969
1970 #[async_trait]
1971 impl FetchProvider for FailingFetch {
1972 async fn fetch(
1973 &self,
1974 url: &str,
1975 _budget: crate::providers::ctx::WebByteBudget,
1976 ) -> Result<WebFetchResult, WebFetchError> {
1977 Err(WebFetchError::HttpStatus {
1978 status: 503,
1979 url: url.to_string(),
1980 })
1981 }
1982 }
1983
1984 let tool = WebFetchTool::new_with_test_snapshots(Arc::new(FailingFetch), "mock");
1985 let (ctx, _rx) =
1986 test_exec_context(TurnId(9), ToolCallId(9), std::path::PathBuf::from("/tmp"));
1987 let outcome = tool
1988 .execute(serde_json::json!({"url": "https://example.com/page"}), ctx)
1989 .await;
1990
1991 assert_eq!(outcome.status, ToolStatus::Error);
1992 match &outcome.metadata.detail {
1993 ToolMetadata::WebFetch {
1994 status,
1995 error_kind,
1996 backend,
1997 final_url,
1998 ..
1999 } => {
2000 assert_eq!(*status, Some(503));
2001 assert_eq!(error_kind.as_deref(), Some("http_status"));
2002 assert_eq!(backend, "mock");
2003 assert!(final_url.is_none());
2004 },
2005 other => panic!("expected web metadata, got {other:?}"),
2006 }
2007 }
2008
2009 #[tokio::test]
2010 async fn web_search_progress_redacts_without_changing_transport_query() {
2011 use crate::providers::ctx::test_exec_context;
2012 use mermaid_domain::{ToolCallId, ToolStatus, TurnId};
2013
2014 struct RecordingSearch {
2015 seen: Arc<Mutex<Option<String>>>,
2016 }
2017
2018 #[async_trait]
2019 impl SearchProvider for RecordingSearch {
2020 async fn search(
2021 &self,
2022 query: &str,
2023 _count: usize,
2024 _budget: crate::providers::ctx::WebByteBudget,
2025 ) -> anyhow::Result<Vec<crate::providers::tool::web_client::SearchResult>> {
2026 *self
2027 .seen
2028 .lock()
2029 .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(query.to_string());
2030 Ok(Vec::new())
2031 }
2032 }
2033
2034 let seen = Arc::new(Mutex::new(None));
2035 let tool = WebSearchTool {
2036 backend: Arc::new(RecordingSearch { seen: seen.clone() }),
2037 backend_name: "mock",
2038 };
2039 let (ctx, mut progress) =
2040 test_exec_context(TurnId(91), ToolCallId(91), std::path::PathBuf::from("/tmp"));
2041 let query = "research OPENAI_API_KEY=abc";
2042 let outcome = tool.execute(serde_json::json!({"query": query}), ctx).await;
2043 assert_eq!(outcome.status, ToolStatus::Success);
2044 assert_eq!(
2045 seen.lock()
2046 .unwrap_or_else(std::sync::PoisonError::into_inner)
2047 .as_deref(),
2048 Some(query),
2049 "redaction must not alter the transport query"
2050 );
2051 let ProgressEvent::Status(status) = progress.recv().await.expect("search progress") else {
2052 panic!("expected search status progress");
2053 };
2054 assert!(
2055 !status.contains("abc"),
2056 "progress leaked query secret: {status}"
2057 );
2058 assert!(status.contains("OPENAI_API_KEY=[REDACTED]"));
2059 }
2060
2061 #[tokio::test]
2062 async fn snapshot_line_ranges_do_not_refetch_mutable_pages() {
2063 use crate::providers::ctx::test_exec_context;
2064 use mermaid_domain::{ToolCallId, ToolStatus, TurnId};
2065 use std::sync::atomic::{AtomicUsize, Ordering};
2066
2067 struct MockFetch {
2068 calls: Arc<AtomicUsize>,
2069 }
2070
2071 #[async_trait]
2072 impl FetchProvider for MockFetch {
2073 async fn fetch(
2074 &self,
2075 url: &str,
2076 _budget: crate::providers::ctx::WebByteBudget,
2077 ) -> Result<WebFetchResult, WebFetchError> {
2078 self.calls.fetch_add(1, Ordering::SeqCst);
2079 let mut result = page("one\ntwo\nthree");
2080 result.requested_url = url.to_string();
2081 result.final_url = Some(url.to_string());
2082 Ok(result)
2083 }
2084 }
2085
2086 let calls = Arc::new(AtomicUsize::new(0));
2087 let tool = WebFetchTool::new_with_test_snapshots(
2088 Arc::new(MockFetch {
2089 calls: calls.clone(),
2090 }),
2091 "mock",
2092 );
2093 let (mut ctx, _rx) =
2094 test_exec_context(TurnId(10), ToolCallId(10), std::path::PathBuf::from("/tmp"));
2095 ctx.session_id = Some("session-a".to_string());
2096 let first = tool
2097 .execute(serde_json::json!({"url": "https://example.com/page"}), ctx)
2098 .await;
2099 assert_eq!(first.status, ToolStatus::Success);
2100 let snapshot_id = match &first.metadata.detail {
2101 ToolMetadata::WebFetch { snapshot_id, .. } => snapshot_id.clone().expect("snapshot id"),
2102 other => panic!("expected web metadata, got {other:?}"),
2103 };
2104
2105 let (mut foreign_ctx, _rx) =
2106 test_exec_context(TurnId(11), ToolCallId(11), std::path::PathBuf::from("/tmp"));
2107 foreign_ctx.session_id = Some("session-b".to_string());
2108 let foreign = tool
2109 .execute(
2110 serde_json::json!({
2111 "snapshot_id": snapshot_id.clone(),
2112 "start_line": 2,
2113 "line_count": 1
2114 }),
2115 foreign_ctx,
2116 )
2117 .await;
2118 assert_eq!(foreign.status, ToolStatus::Error);
2119 assert!(foreign.output().contains("unavailable or was evicted"));
2120
2121 let (mut ctx, _rx) =
2122 test_exec_context(TurnId(12), ToolCallId(12), std::path::PathBuf::from("/tmp"));
2123 ctx.session_id = Some("session-a".to_string());
2124 let continuation = tool
2125 .execute(
2126 serde_json::json!({
2127 "snapshot_id": snapshot_id,
2128 "start_line": 2,
2129 "line_count": 1
2130 }),
2131 ctx,
2132 )
2133 .await;
2134 assert_eq!(continuation.status, ToolStatus::Success);
2135 assert!(continuation.output().contains("L2: two"));
2136 assert!(!continuation.output().contains("L1: one"));
2137 assert_eq!(
2138 calls.load(Ordering::SeqCst),
2139 1,
2140 "snapshot triggered a refetch"
2141 );
2142 }
2143
2144 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2145 async fn snapshot_blocking_work_respects_the_global_extractor_limit() {
2146 use std::sync::atomic::{AtomicUsize, Ordering};
2147
2148 let active = Arc::new(AtomicUsize::new(0));
2149 let peak = Arc::new(AtomicUsize::new(0));
2150 let mut jobs = Vec::new();
2151 for _ in 0..8 {
2152 let active = Arc::clone(&active);
2153 let peak = Arc::clone(&peak);
2154 jobs.push(tokio::spawn(async move {
2155 run_snapshot_blocking(move || {
2156 let now = active.fetch_add(1, Ordering::SeqCst) + 1;
2157 peak.fetch_max(now, Ordering::SeqCst);
2158 std::thread::sleep(std::time::Duration::from_millis(20));
2159 active.fetch_sub(1, Ordering::SeqCst);
2160 })
2161 .await
2162 .unwrap();
2163 }));
2164 }
2165 for job in jobs {
2166 job.await.unwrap();
2167 }
2168
2169 assert_eq!(active.load(Ordering::SeqCst), 0);
2170 assert!(
2171 peak.load(Ordering::SeqCst) <= mermaid_model::constants::MAX_WEB_EXTRACTION_CONCURRENCY
2172 );
2173 }
2174
2175 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2176 async fn cancelled_snapshot_waiter_keeps_its_permit_until_blocking_work_finishes() {
2177 let limiter = Arc::new(tokio::sync::Semaphore::new(1));
2178 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
2179 let (release_tx, release_rx) = std::sync::mpsc::sync_channel(0);
2180 let worker_limiter = limiter.clone();
2181 let worker = tokio::spawn(async move {
2182 run_snapshot_blocking_with(worker_limiter, move || {
2183 let _ = started_tx.send(());
2184 release_rx.recv().expect("test releases blocking worker");
2185 })
2186 .await
2187 });
2188 started_rx.await.expect("blocking worker started");
2189 worker.abort();
2190 let _ = worker.await;
2191 assert_eq!(
2192 limiter.available_permits(),
2193 0,
2194 "cancelling the async waiter released a still-running blocking job"
2195 );
2196
2197 release_tx.send(()).expect("release blocking worker");
2198 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2199 while limiter.available_permits() == 0 {
2200 tokio::task::yield_now().await;
2201 }
2202 })
2203 .await
2204 .expect("blocking worker did not release its permit");
2205 }
2206
2207 #[tokio::test]
2208 async fn web_search_batch_survives_empty_and_failed_queries() {
2209 use crate::providers::ctx::test_exec_context;
2210 use crate::providers::tool::web_client::SearchResult;
2211 use async_trait::async_trait;
2212 use mermaid_domain::{ToolCallId, ToolStatus, TurnId};
2213 use std::sync::Arc;
2214
2215 struct Mock;
2216 #[async_trait]
2217 impl SearchProvider for Mock {
2218 async fn search(
2219 &self,
2220 query: &str,
2221 _count: usize,
2222 _budget: crate::providers::ctx::WebByteBudget,
2223 ) -> anyhow::Result<Vec<SearchResult>> {
2224 match query {
2225 "boom" => Err(anyhow::anyhow!("backend down")),
2226 "empty" => Ok(Vec::new()),
2227 _ => Ok(vec![SearchResult {
2228 title: "Title".to_string(),
2229 url: "https://example.com".to_string(),
2230 snippet: "snip".to_string(),
2231 full_content: "content".to_string(),
2232 }]),
2233 }
2234 }
2235 }
2236
2237 let mk = || WebSearchTool {
2238 backend: Arc::new(Mock),
2239 backend_name: "mock",
2240 };
2241 let tmp = std::path::PathBuf::from("/tmp");
2242
2243 let (ctx, _rx) = test_exec_context(TurnId(1), ToolCallId(1), tmp.clone());
2245 let out = mk()
2246 .execute(
2247 serde_json::json!({"queries": [{"query":"good"},{"query":"empty"},{"query":"boom"}]}),
2248 ctx,
2249 )
2250 .await;
2251 assert_eq!(
2252 out.status,
2253 ToolStatus::Success,
2254 "a partial batch must not abort"
2255 );
2256 assert!(
2257 out.output().contains("https://example.com"),
2258 "keeps the good result"
2259 );
2260 match &out.metadata.detail {
2261 ToolMetadata::WebSearch {
2262 partial, failures, ..
2263 } => {
2264 assert!(*partial);
2265 assert_eq!(failures.len(), 1);
2266 assert_eq!(failures[0].query_index, 2);
2267 assert!(failures[0].error.contains("backend down"));
2268 },
2269 other => panic!("expected web search metadata, got {other:?}"),
2270 }
2271
2272 let (ctx, _rx) = test_exec_context(TurnId(2), ToolCallId(2), tmp.clone());
2274 let out = mk()
2275 .execute(serde_json::json!({"query": "empty"}), ctx)
2276 .await;
2277 assert_eq!(out.status, ToolStatus::Success, "empty is not an error");
2278 assert!(out.output().contains("no results"));
2279
2280 let (ctx, _rx) = test_exec_context(TurnId(3), ToolCallId(3), tmp);
2282 let out = mk()
2283 .execute(
2284 serde_json::json!({"queries": [{"query":"boom"},{"query":"boom"}]}),
2285 ctx,
2286 )
2287 .await;
2288 assert_eq!(out.status, ToolStatus::Error, "total failure is an error");
2289 match &out.metadata.detail {
2290 ToolMetadata::WebSearch {
2291 failed_queries,
2292 failures,
2293 ..
2294 } => {
2295 assert_eq!(*failed_queries, 2);
2296 assert_eq!(failures.len(), 2);
2297 },
2298 other => panic!("expected web search metadata, got {other:?}"),
2299 }
2300 }
2301
2302 #[tokio::test]
2303 async fn web_search_batch_caps_concurrency_and_restores_input_order() {
2304 use crate::providers::ctx::test_exec_context;
2305 use crate::providers::tool::web_client::SearchResult;
2306 use mermaid_domain::{ToolCallId, ToolStatus, TurnId};
2307 use std::sync::atomic::{AtomicUsize, Ordering};
2308
2309 struct ConcurrencyMock {
2310 active: Arc<AtomicUsize>,
2311 peak: Arc<AtomicUsize>,
2312 }
2313
2314 #[async_trait]
2315 impl SearchProvider for ConcurrencyMock {
2316 async fn search(
2317 &self,
2318 query: &str,
2319 _count: usize,
2320 _budget: crate::providers::ctx::WebByteBudget,
2321 ) -> anyhow::Result<Vec<SearchResult>> {
2322 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
2323 self.peak.fetch_max(active, Ordering::SeqCst);
2324 let delay = 5 + (6 - query.parse::<u64>().unwrap_or(0)) * 5;
2325 tokio::time::sleep(std::time::Duration::from_millis(delay)).await;
2326 self.active.fetch_sub(1, Ordering::SeqCst);
2327 Ok(vec![SearchResult {
2328 title: format!("result {query}"),
2329 url: format!("https://example.com/{query}"),
2330 snippet: String::new(),
2331 full_content: format!("content {query}"),
2332 }])
2333 }
2334 }
2335
2336 let active = Arc::new(AtomicUsize::new(0));
2337 let peak = Arc::new(AtomicUsize::new(0));
2338 let tool = WebSearchTool {
2339 backend: Arc::new(ConcurrencyMock {
2340 active: active.clone(),
2341 peak: peak.clone(),
2342 }),
2343 backend_name: "mock",
2344 };
2345 let (ctx, _rx) =
2346 test_exec_context(TurnId(20), ToolCallId(20), std::path::PathBuf::from("/tmp"));
2347 let queries: Vec<_> = (0..6)
2348 .map(|index| serde_json::json!({"query": index.to_string()}))
2349 .collect();
2350 let outcome = tool
2351 .execute(serde_json::json!({"queries": queries}), ctx)
2352 .await;
2353 assert_eq!(outcome.status, ToolStatus::Success);
2354 assert_eq!(active.load(Ordering::SeqCst), 0);
2355 assert_eq!(
2356 peak.load(Ordering::SeqCst),
2357 mermaid_model::constants::MAX_WEB_SEARCH_CONCURRENCY
2358 );
2359 let mut cursor = 0;
2360 for index in 0..6 {
2361 let marker = format!("=== query: {index} ===");
2362 let position = outcome.output()[cursor..]
2363 .find(&marker)
2364 .map(|offset| cursor + offset)
2365 .expect("ordered query section");
2366 assert!(position >= cursor);
2367 cursor = position + marker.len();
2368 }
2369 }
2370
2371 #[tokio::test]
2372 async fn web_search_complete_output_budget_is_byte_exact_for_multibyte_text() {
2373 use crate::providers::ctx::test_exec_context;
2374 use crate::providers::tool::web_client::SearchResult;
2375 use mermaid_domain::{ToolCallId, ToolStatus, TurnId};
2376
2377 struct MultibyteMock;
2378
2379 #[async_trait]
2380 impl SearchProvider for MultibyteMock {
2381 async fn search(
2382 &self,
2383 _query: &str,
2384 _count: usize,
2385 _budget: crate::providers::ctx::WebByteBudget,
2386 ) -> anyhow::Result<Vec<SearchResult>> {
2387 Ok(vec![SearchResult {
2388 title: "界".repeat(5_000),
2389 url: "https://example.com/result".to_string(),
2390 snippet: String::new(),
2391 full_content: "界".repeat(20_000),
2392 }])
2393 }
2394 }
2395
2396 let tool = WebSearchTool {
2397 backend: Arc::new(MultibyteMock),
2398 backend_name: "mock",
2399 };
2400 let (ctx, _rx) =
2401 test_exec_context(TurnId(21), ToolCallId(21), std::path::PathBuf::from("/tmp"));
2402 let outcome = tool
2403 .execute(serde_json::json!({"query": "multibyte"}), ctx)
2404 .await;
2405
2406 assert_eq!(outcome.status, ToolStatus::Success);
2407 assert!(
2408 outcome.output().len() <= mermaid_model::constants::WEB_SEARCH_AGGREGATE_MAX_BYTES,
2409 "{} bytes escaped the complete search-result cap",
2410 outcome.output().len()
2411 );
2412 assert!(std::str::from_utf8(outcome.output().as_bytes()).is_ok());
2413 }
2414
2415 #[tokio::test]
2416 async fn web_search_total_failure_and_structured_errors_are_byte_bounded() {
2417 use crate::providers::ctx::test_exec_context;
2418 use mermaid_domain::{ToolCallId, ToolStatus, TurnId};
2419
2420 struct LargeFailure;
2421
2422 #[async_trait]
2423 impl SearchProvider for LargeFailure {
2424 async fn search(
2425 &self,
2426 _query: &str,
2427 _count: usize,
2428 _budget: crate::providers::ctx::WebByteBudget,
2429 ) -> anyhow::Result<Vec<crate::providers::tool::web_client::SearchResult>> {
2430 Err(anyhow::anyhow!("{}", "界".repeat(20_000)))
2431 }
2432 }
2433
2434 let tool = WebSearchTool {
2435 backend: Arc::new(LargeFailure),
2436 backend_name: "mock",
2437 };
2438 let (ctx, _rx) =
2439 test_exec_context(TurnId(22), ToolCallId(22), std::path::PathBuf::from("/tmp"));
2440 let outcome = tool
2441 .execute(
2442 serde_json::json!({"queries": [{"query": "one"}, {"query": "two"}]}),
2443 ctx,
2444 )
2445 .await;
2446
2447 assert_eq!(outcome.status, ToolStatus::Error);
2448 assert!(
2449 outcome.output().len() <= mermaid_model::constants::WEB_SEARCH_AGGREGATE_MAX_BYTES,
2450 "{} bytes escaped the complete search error cap",
2451 outcome.output().len()
2452 );
2453 let ToolMetadata::WebSearch { failures, .. } = &outcome.metadata.detail else {
2454 panic!("expected web search metadata");
2455 };
2456 assert_eq!(failures.len(), 2);
2457 assert!(
2458 failures
2459 .iter()
2460 .all(|failure| failure.error.len() <= MAX_WEB_SEARCH_FAILURE_BYTES)
2461 );
2462 }
2463}