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