| 1 | //! Search backend selection and the shared async adapter contract. |
| 2 | |
| 3 | use std::time::{Duration, Instant}; |
| 4 | |
| 5 | use async_trait::async_trait; |
| 6 | |
| 7 | use super::contract::{BackendId, BackendSearch, DegradedReason, QueryCapabilities, SearchQuery}; |
| 8 | use super::contract::{CapabilityState as QueryCapabilityState, SearchResult}; |
| 9 | use crate::client::ProviderNativeSearchRequest; |
| 10 | use crate::config::SearchProvider; |
| 11 | use crate::tools::spec::{ToolContext, ToolError}; |
| 12 | |
| 13 | #[async_trait] |
| 14 | pub(crate) trait SearchBackend: Send + Sync { |
| 15 | fn id(&self) -> BackendId; |
| 16 | fn capabilities(&self) -> QueryCapabilities; |
| 17 | async fn search( |
| 18 | &self, |
| 19 | query: &SearchQuery, |
| 20 | deadline: Instant, |
| 21 | ) -> Result<BackendSearch, ToolError>; |
| 22 | } |
| 23 | |
| 24 | #[derive(Clone, Copy)] |
| 25 | pub(crate) struct BackendContext<'a> { |
| 26 | tool_context: &'a ToolContext, |
| 27 | } |
| 28 | |
| 29 | pub(crate) enum ConfiguredSearchBackend<'a> { |
| 30 | Bing(BackendContext<'a>), |
| 31 | DuckDuckGo(BackendContext<'a>), |
| 32 | Tavily(BackendContext<'a>), |
| 33 | Bocha(BackendContext<'a>), |
| 34 | Metaso(BackendContext<'a>), |
| 35 | Searxng(BackendContext<'a>), |
| 36 | Baidu(BackendContext<'a>), |
| 37 | Volcengine(BackendContext<'a>), |
| 38 | Sofya(BackendContext<'a>), |
| 39 | } |
| 40 | |
| 41 | #[derive(Clone, Copy)] |
| 42 | struct ProviderNativeSearchBackend<'a> { |
| 43 | context: &'a ToolContext, |
| 44 | } |
| 45 | |
| 46 | impl<'a> ConfiguredSearchBackend<'a> { |
| 47 | #[must_use] |
| 48 | pub(crate) fn from_provider(context: &'a ToolContext, provider: SearchProvider) -> Self { |
| 49 | let backend = BackendContext { |
| 50 | tool_context: context, |
| 51 | }; |
| 52 | match provider { |
| 53 | SearchProvider::Bing => Self::Bing(backend), |
| 54 | SearchProvider::DuckDuckGo => Self::DuckDuckGo(backend), |
| 55 | SearchProvider::Tavily => Self::Tavily(backend), |
| 56 | SearchProvider::Bocha => Self::Bocha(backend), |
| 57 | SearchProvider::Metaso => Self::Metaso(backend), |
| 58 | SearchProvider::Searxng => Self::Searxng(backend), |
| 59 | SearchProvider::Baidu => Self::Baidu(backend), |
| 60 | SearchProvider::Volcengine => Self::Volcengine(backend), |
| 61 | SearchProvider::Sofya => Self::Sofya(backend), |
| 62 | } |
| 63 | } |
| 64 | |
| 65 | const fn provider(&self) -> SearchProvider { |
| 66 | match self { |
| 67 | Self::Bing(_) => SearchProvider::Bing, |
| 68 | Self::DuckDuckGo(_) => SearchProvider::DuckDuckGo, |
| 69 | Self::Tavily(_) => SearchProvider::Tavily, |
| 70 | Self::Bocha(_) => SearchProvider::Bocha, |
| 71 | Self::Metaso(_) => SearchProvider::Metaso, |
| 72 | Self::Searxng(_) => SearchProvider::Searxng, |
| 73 | Self::Baidu(_) => SearchProvider::Baidu, |
| 74 | Self::Volcengine(_) => SearchProvider::Volcengine, |
| 75 | Self::Sofya(_) => SearchProvider::Sofya, |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | const fn context(&self) -> &BackendContext<'a> { |
| 80 | match self { |
| 81 | Self::Bing(context) |
| 82 | | Self::DuckDuckGo(context) |
| 83 | | Self::Tavily(context) |
| 84 | | Self::Bocha(context) |
| 85 | | Self::Metaso(context) |
| 86 | | Self::Searxng(context) |
| 87 | | Self::Baidu(context) |
| 88 | | Self::Volcengine(context) |
| 89 | | Self::Sofya(context) => context, |
| 90 | } |
| 91 | } |
| 92 | } |
| 93 | |
| 94 | pub(crate) struct SearchBackendChain<'a> { |
| 95 | backends: Vec<Box<dyn SearchBackend + 'a>>, |
| 96 | } |
| 97 | |
| 98 | #[derive(Debug)] |
| 99 | pub(crate) struct ChainedSearch { |
| 100 | pub(crate) raw: BackendSearch, |
| 101 | pub(crate) capabilities: QueryCapabilities, |
| 102 | } |
| 103 | |
| 104 | impl<'a> SearchBackendChain<'a> { |
| 105 | #[must_use] |
| 106 | pub(crate) fn from_context(context: &'a ToolContext) -> Self { |
| 107 | let selected = context.search_provider; |
| 108 | let mut backends: Vec<Box<dyn SearchBackend + 'a>> = Vec::new(); |
| 109 | if should_prepend_provider_native(context) { |
| 110 | backends.push(Box::new(ProviderNativeSearchBackend { context })); |
| 111 | } |
| 112 | backends.push(Box::new(ConfiguredSearchBackend::from_provider( |
| 113 | context, selected, |
| 114 | ))); |
| 115 | if !matches!(selected, SearchProvider::Bing | SearchProvider::DuckDuckGo) { |
| 116 | backends.push(Box::new(ConfiguredSearchBackend::from_provider( |
| 117 | context, |
| 118 | SearchProvider::DuckDuckGo, |
| 119 | ))); |
| 120 | } |
| 121 | Self { backends } |
| 122 | } |
| 123 | |
| 124 | #[must_use] |
| 125 | pub(crate) fn initial_backend(&self) -> BackendId { |
| 126 | self.backends |
| 127 | .first() |
| 128 | .expect("a search chain always has a configured backend") |
| 129 | .id() |
| 130 | } |
| 131 | |
| 132 | pub(crate) async fn search( |
| 133 | &self, |
| 134 | query: &SearchQuery, |
| 135 | deadline: Instant, |
| 136 | first_attempt_budget: Option<Duration>, |
| 137 | ) -> Result<ChainedSearch, ToolError> { |
| 138 | let backends = self |
| 139 | .backends |
| 140 | .iter() |
| 141 | .map(|backend| backend.as_ref()) |
| 142 | .collect::<Vec<_>>(); |
| 143 | run_backend_chain(&backends, query, deadline, first_attempt_budget).await |
| 144 | } |
| 145 | } |
| 146 | |
| 147 | fn should_prepend_provider_native(context: &ToolContext) -> bool { |
| 148 | provider_native_is_available( |
| 149 | context |
| 150 | .route_capabilities |
| 151 | .server_side_web_search |
| 152 | .is_supported(), |
| 153 | context.provider_native_search.is_some(), |
| 154 | ) |
| 155 | } |
| 156 | |
| 157 | const fn provider_native_is_available(capability_supported: bool, client_present: bool) -> bool { |
| 158 | capability_supported && client_present |
| 159 | } |
| 160 | |
| 161 | async fn run_backend_chain( |
| 162 | backends: &[&dyn SearchBackend], |
| 163 | query: &SearchQuery, |
| 164 | deadline: Instant, |
| 165 | first_attempt_budget: Option<Duration>, |
| 166 | ) -> Result<ChainedSearch, ToolError> { |
| 167 | let mut degraded = Vec::new(); |
| 168 | let mut last_empty = None; |
| 169 | let mut attempted = Vec::new(); |
| 170 | |
| 171 | for (index, backend) in backends.iter().enumerate() { |
| 172 | let remaining = deadline.saturating_duration_since(Instant::now()); |
| 173 | if remaining.is_zero() { |
| 174 | break; |
| 175 | } |
| 176 | let backend_id = backend.id(); |
| 177 | if let Some(previous) = attempted.last() { |
| 178 | degraded.push(DegradedReason::BackendFallback { |
| 179 | from: *previous, |
| 180 | to: backend_id, |
| 181 | }); |
| 182 | } |
| 183 | attempted.push(backend_id); |
| 184 | |
| 185 | let attempts_left = u32::try_from(backends.len() - index).unwrap_or(u32::MAX); |
| 186 | let fair_share = remaining / attempts_left; |
| 187 | let attempt_budget = if index == 0 { |
| 188 | first_attempt_budget |
| 189 | .map(|budget| budget.min(remaining)) |
| 190 | .unwrap_or(fair_share) |
| 191 | } else { |
| 192 | fair_share |
| 193 | } |
| 194 | .max(Duration::from_millis(1)); |
| 195 | let attempt_deadline = Instant::now() + attempt_budget; |
| 196 | |
| 197 | let result = tokio::time::timeout(attempt_budget, backend.search(query, attempt_deadline)) |
| 198 | .await |
| 199 | .map_err(|_| ToolError::Timeout { |
| 200 | seconds: u64::try_from(attempt_budget.as_millis()) |
| 201 | .unwrap_or(u64::MAX) |
| 202 | .div_ceil(1_000), |
| 203 | }) |
| 204 | .and_then(std::convert::identity); |
| 205 | |
| 206 | match result { |
| 207 | Ok(mut raw) if !raw.results.is_empty() => { |
| 208 | degraded.append(&mut raw.degraded); |
| 209 | raw.degraded = degraded; |
| 210 | return Ok(ChainedSearch { |
| 211 | raw, |
| 212 | capabilities: backend.capabilities(), |
| 213 | }); |
| 214 | } |
| 215 | Ok(mut raw) => { |
| 216 | degraded.push(DegradedReason::NoUsableResults { |
| 217 | backend: backend_id, |
| 218 | }); |
| 219 | degraded.append(&mut raw.degraded); |
| 220 | last_empty = Some((raw, backend.capabilities())); |
| 221 | } |
| 222 | Err(error) if is_fail_closed(&error) => return Err(error), |
| 223 | Err(error) if backends.len() == 1 => return Err(error), |
| 224 | Err(_) => degraded.push(DegradedReason::BackendUnavailable { |
| 225 | backend: backend_id, |
| 226 | }), |
| 227 | } |
| 228 | } |
| 229 | |
| 230 | if let Some((mut raw, capabilities)) = last_empty { |
| 231 | raw.degraded = degraded; |
| 232 | return Ok(ChainedSearch { raw, capabilities }); |
| 233 | } |
| 234 | |
| 235 | if attempted.is_empty() { |
| 236 | return Err(ToolError::Timeout { seconds: 1 }); |
| 237 | } |
| 238 | |
| 239 | let backend_ids = attempted |
| 240 | .into_iter() |
| 241 | .map(BackendId::as_str) |
| 242 | .collect::<Vec<_>>() |
| 243 | .join(", "); |
| 244 | Err(ToolError::not_available(format!( |
| 245 | "web search backends unavailable: {backend_ids}" |
| 246 | ))) |
| 247 | } |
| 248 | |
| 249 | const fn is_fail_closed(error: &ToolError) -> bool { |
| 250 | matches!( |
| 251 | error, |
| 252 | ToolError::InvalidInput { .. } |
| 253 | | ToolError::MissingField { .. } |
| 254 | | ToolError::PathEscape { .. } |
| 255 | | ToolError::Cancelled { .. } |
| 256 | | ToolError::PermissionDenied { .. } |
| 257 | ) |
| 258 | } |
| 259 | |
| 260 | #[async_trait] |
| 261 | impl SearchBackend for ConfiguredSearchBackend<'_> { |
| 262 | fn id(&self) -> BackendId { |
| 263 | match self.provider() { |
| 264 | SearchProvider::Bing => BackendId::Bing, |
| 265 | SearchProvider::DuckDuckGo => BackendId::DuckDuckGo, |
| 266 | SearchProvider::Tavily => BackendId::Tavily, |
| 267 | SearchProvider::Bocha => BackendId::Bocha, |
| 268 | SearchProvider::Metaso => BackendId::Metaso, |
| 269 | SearchProvider::Searxng => BackendId::Searxng, |
| 270 | SearchProvider::Baidu => BackendId::Baidu, |
| 271 | SearchProvider::Volcengine => BackendId::Volcengine, |
| 272 | SearchProvider::Sofya => BackendId::Sofya, |
| 273 | } |
| 274 | } |
| 275 | |
| 276 | fn capabilities(&self) -> QueryCapabilities { |
| 277 | // All current adapters enforce result count. Other knobs are either |
| 278 | // post-filtered by the shared harness or reported as not honored. |
| 279 | QueryCapabilities::count_only() |
| 280 | } |
| 281 | |
| 282 | async fn search( |
| 283 | &self, |
| 284 | query: &SearchQuery, |
| 285 | deadline: Instant, |
| 286 | ) -> Result<BackendSearch, ToolError> { |
| 287 | crate::tools::web_search::run_backend_search( |
| 288 | self.provider(), |
| 289 | query, |
| 290 | deadline, |
| 291 | self.context().tool_context, |
| 292 | ) |
| 293 | .await |
| 294 | } |
| 295 | } |
| 296 | |
| 297 | #[async_trait] |
| 298 | impl SearchBackend for ProviderNativeSearchBackend<'_> { |
| 299 | fn id(&self) -> BackendId { |
| 300 | BackendId::ProviderNative |
| 301 | } |
| 302 | |
| 303 | fn capabilities(&self) -> QueryCapabilities { |
| 304 | QueryCapabilities { |
| 305 | max_results: QueryCapabilityState::Supported, |
| 306 | recency: QueryCapabilityState::Unsupported, |
| 307 | domains: QueryCapabilityState::Supported, |
| 308 | locale: QueryCapabilityState::Unsupported, |
| 309 | published_date: QueryCapabilityState::Unknown, |
| 310 | } |
| 311 | } |
| 312 | |
| 313 | async fn search( |
| 314 | &self, |
| 315 | query: &SearchQuery, |
| 316 | _deadline: Instant, |
| 317 | ) -> Result<BackendSearch, ToolError> { |
| 318 | if !self |
| 319 | .context |
| 320 | .route_capabilities |
| 321 | .server_side_web_search |
| 322 | .is_supported() |
| 323 | { |
| 324 | return Err(ToolError::not_available( |
| 325 | "active route does not report provider-native web search", |
| 326 | )); |
| 327 | } |
| 328 | let client = self |
| 329 | .context |
| 330 | .provider_native_search |
| 331 | .as_ref() |
| 332 | .ok_or_else(|| ToolError::not_available("provider-native search client unavailable"))?; |
| 333 | if let Some(maximum) = client.maximum_domain_count() |
| 334 | && query.domains.len() > maximum |
| 335 | { |
| 336 | return Err(ToolError::invalid_input(format!( |
| 337 | "{} native web search accepts at most {maximum} domains", |
| 338 | client.provider().as_str() |
| 339 | ))); |
| 340 | } |
| 341 | let host = client.host().ok_or_else(|| { |
| 342 | ToolError::execution_failed("provider-native search endpoint has no valid host") |
| 343 | })?; |
| 344 | crate::tools::web_search::check_policy( |
| 345 | self.context.network_policy.as_ref(), |
| 346 | host.as_str(), |
| 347 | )?; |
| 348 | let response = client |
| 349 | .search(&ProviderNativeSearchRequest { |
| 350 | query: query.query.clone(), |
| 351 | max_results: query.max_results, |
| 352 | domains: query.domains.clone(), |
| 353 | }) |
| 354 | .await |
| 355 | .map_err(|error| { |
| 356 | ToolError::execution_failed(format!( |
| 357 | "{} provider-native web search failed: {error}", |
| 358 | client.provider().as_str() |
| 359 | )) |
| 360 | })?; |
| 361 | let results = response |
| 362 | .citations |
| 363 | .into_iter() |
| 364 | .enumerate() |
| 365 | .map(|(index, citation)| { |
| 366 | SearchResult::new( |
| 367 | index + 1, |
| 368 | citation.title, |
| 369 | citation.url, |
| 370 | citation.snippet, |
| 371 | citation.published, |
| 372 | ) |
| 373 | }) |
| 374 | .collect(); |
| 375 | Ok(BackendSearch { |
| 376 | backend: BackendId::ProviderNative, |
| 377 | source: format!( |
| 378 | "provider-native/{}/{}", |
| 379 | client.provider().as_str(), |
| 380 | client.model() |
| 381 | ), |
| 382 | backend_detail: Some(host), |
| 383 | results, |
| 384 | degraded: Vec::new(), |
| 385 | note: response.answer, |
| 386 | }) |
| 387 | } |
| 388 | } |
| 389 | |
| 390 | #[cfg(test)] |
| 391 | mod tests { |
| 392 | use std::sync::{Arc, Mutex}; |
| 393 | |
| 394 | use super::*; |
| 395 | |
| 396 | struct FakeBackend { |
| 397 | id: BackendId, |
| 398 | result: Result<Vec<super::super::contract::SearchResult>, ToolError>, |
| 399 | } |
| 400 | |
| 401 | #[async_trait] |
| 402 | impl SearchBackend for FakeBackend { |
| 403 | fn id(&self) -> BackendId { |
| 404 | self.id |
| 405 | } |
| 406 | |
| 407 | fn capabilities(&self) -> QueryCapabilities { |
| 408 | QueryCapabilities::count_only() |
| 409 | } |
| 410 | |
| 411 | async fn search( |
| 412 | &self, |
| 413 | _query: &SearchQuery, |
| 414 | _deadline: Instant, |
| 415 | ) -> Result<BackendSearch, ToolError> { |
| 416 | Ok(BackendSearch { |
| 417 | backend: self.id, |
| 418 | source: self.id.as_str().to_string(), |
| 419 | backend_detail: None, |
| 420 | results: self.result.clone()?, |
| 421 | degraded: Vec::new(), |
| 422 | note: None, |
| 423 | }) |
| 424 | } |
| 425 | } |
| 426 | |
| 427 | fn query() -> SearchQuery { |
| 428 | SearchQuery::new("bounded chain".to_string(), 5, None, Vec::new(), None) |
| 429 | } |
| 430 | |
| 431 | fn result() -> super::super::contract::SearchResult { |
| 432 | super::super::contract::SearchResult::new( |
| 433 | 1, |
| 434 | "Fallback result".to_string(), |
| 435 | "https://example.com/result".to_string(), |
| 436 | None, |
| 437 | None, |
| 438 | ) |
| 439 | } |
| 440 | |
| 441 | #[test] |
| 442 | fn every_configured_provider_maps_to_one_explicit_backend_adapter() { |
| 443 | let cases = [ |
| 444 | (SearchProvider::Bing, BackendId::Bing), |
| 445 | (SearchProvider::DuckDuckGo, BackendId::DuckDuckGo), |
| 446 | (SearchProvider::Tavily, BackendId::Tavily), |
| 447 | (SearchProvider::Bocha, BackendId::Bocha), |
| 448 | (SearchProvider::Metaso, BackendId::Metaso), |
| 449 | (SearchProvider::Searxng, BackendId::Searxng), |
| 450 | (SearchProvider::Baidu, BackendId::Baidu), |
| 451 | (SearchProvider::Volcengine, BackendId::Volcengine), |
| 452 | (SearchProvider::Sofya, BackendId::Sofya), |
| 453 | ]; |
| 454 | |
| 455 | for (provider, expected) in cases { |
| 456 | let mut context = ToolContext::new(std::path::PathBuf::from(".")); |
| 457 | context.search_provider = provider; |
| 458 | let backend = ConfiguredSearchBackend::from_provider(&context, provider); |
| 459 | assert_eq!(backend.id(), expected); |
| 460 | assert_eq!( |
| 461 | backend.capabilities().max_results, |
| 462 | super::super::contract::CapabilityState::Supported |
| 463 | ); |
| 464 | } |
| 465 | } |
| 466 | |
| 467 | #[test] |
| 468 | fn provider_native_is_fail_closed_without_both_fact_and_client() { |
| 469 | assert!(!provider_native_is_available(false, false)); |
| 470 | assert!(!provider_native_is_available(true, false)); |
| 471 | assert!(!provider_native_is_available(false, true)); |
| 472 | assert!(provider_native_is_available(true, true)); |
| 473 | } |
| 474 | |
| 475 | #[tokio::test] |
| 476 | async fn unavailable_api_falls_back_with_explicit_receipts() { |
| 477 | let api = FakeBackend { |
| 478 | id: BackendId::Tavily, |
| 479 | result: Err(ToolError::execution_failed( |
| 480 | "provider detail must stay private", |
| 481 | )), |
| 482 | }; |
| 483 | let scrape = FakeBackend { |
| 484 | id: BackendId::DuckDuckGo, |
| 485 | result: Ok(vec![result()]), |
| 486 | }; |
| 487 | let response = run_backend_chain( |
| 488 | &[&api, &scrape], |
| 489 | &query(), |
| 490 | Instant::now() + Duration::from_secs(1), |
| 491 | None, |
| 492 | ) |
| 493 | .await |
| 494 | .expect("fallback should succeed"); |
| 495 | |
| 496 | assert_eq!(response.raw.backend, BackendId::DuckDuckGo); |
| 497 | assert_eq!( |
| 498 | response.raw.degraded, |
| 499 | vec![ |
| 500 | DegradedReason::BackendUnavailable { |
| 501 | backend: BackendId::Tavily, |
| 502 | }, |
| 503 | DegradedReason::BackendFallback { |
| 504 | from: BackendId::Tavily, |
| 505 | to: BackendId::DuckDuckGo, |
| 506 | }, |
| 507 | ] |
| 508 | ); |
| 509 | } |
| 510 | |
| 511 | #[tokio::test] |
| 512 | async fn provider_native_to_api_to_scrape_records_every_transition() { |
| 513 | let native = FakeBackend { |
| 514 | id: BackendId::ProviderNative, |
| 515 | result: Err(ToolError::execution_failed("native unavailable")), |
| 516 | }; |
| 517 | let api = FakeBackend { |
| 518 | id: BackendId::Tavily, |
| 519 | result: Err(ToolError::execution_failed("API unavailable")), |
| 520 | }; |
| 521 | let scrape = FakeBackend { |
| 522 | id: BackendId::DuckDuckGo, |
| 523 | result: Ok(vec![result()]), |
| 524 | }; |
| 525 | |
| 526 | let response = run_backend_chain( |
| 527 | &[&native, &api, &scrape], |
| 528 | &query(), |
| 529 | Instant::now() + Duration::from_secs(1), |
| 530 | None, |
| 531 | ) |
| 532 | .await |
| 533 | .expect("final scrape fallback should succeed"); |
| 534 | |
| 535 | assert_eq!(response.raw.backend, BackendId::DuckDuckGo); |
| 536 | assert_eq!( |
| 537 | response.raw.degraded, |
| 538 | vec![ |
| 539 | DegradedReason::BackendUnavailable { |
| 540 | backend: BackendId::ProviderNative, |
| 541 | }, |
| 542 | DegradedReason::BackendFallback { |
| 543 | from: BackendId::ProviderNative, |
| 544 | to: BackendId::Tavily, |
| 545 | }, |
| 546 | DegradedReason::BackendUnavailable { |
| 547 | backend: BackendId::Tavily, |
| 548 | }, |
| 549 | DegradedReason::BackendFallback { |
| 550 | from: BackendId::Tavily, |
| 551 | to: BackendId::DuckDuckGo, |
| 552 | }, |
| 553 | ] |
| 554 | ); |
| 555 | } |
| 556 | |
| 557 | #[tokio::test] |
| 558 | async fn first_attempt_budget_overrides_the_default_fair_share() { |
| 559 | struct DeadlineBackend { |
| 560 | observed_budget: Arc<Mutex<Option<Duration>>>, |
| 561 | } |
| 562 | |
| 563 | #[async_trait] |
| 564 | impl SearchBackend for DeadlineBackend { |
| 565 | fn id(&self) -> BackendId { |
| 566 | BackendId::Volcengine |
| 567 | } |
| 568 | |
| 569 | fn capabilities(&self) -> QueryCapabilities { |
| 570 | QueryCapabilities::count_only() |
| 571 | } |
| 572 | |
| 573 | async fn search( |
| 574 | &self, |
| 575 | _query: &SearchQuery, |
| 576 | deadline: Instant, |
| 577 | ) -> Result<BackendSearch, ToolError> { |
| 578 | *self.observed_budget.lock().expect("budget lock") = |
| 579 | Some(deadline.saturating_duration_since(Instant::now())); |
| 580 | Ok(BackendSearch { |
| 581 | backend: BackendId::Volcengine, |
| 582 | source: "volcengine".to_string(), |
| 583 | backend_detail: None, |
| 584 | results: vec![result()], |
| 585 | degraded: Vec::new(), |
| 586 | note: None, |
| 587 | }) |
| 588 | } |
| 589 | } |
| 590 | |
| 591 | let observed_budget = Arc::new(Mutex::new(None)); |
| 592 | let volcengine = DeadlineBackend { |
| 593 | observed_budget: Arc::clone(&observed_budget), |
| 594 | }; |
| 595 | let fallback = FakeBackend { |
| 596 | id: BackendId::DuckDuckGo, |
| 597 | result: Ok(vec![result()]), |
| 598 | }; |
| 599 | let first_attempt_budget = Duration::from_millis(1_500); |
| 600 | let response = run_backend_chain( |
| 601 | &[&volcengine, &fallback], |
| 602 | &query(), |
| 603 | Instant::now() + Duration::from_secs(2), |
| 604 | Some(first_attempt_budget), |
| 605 | ) |
| 606 | .await |
| 607 | .expect("the first backend should complete inside its dedicated budget"); |
| 608 | |
| 609 | assert_eq!(response.raw.backend, BackendId::Volcengine); |
| 610 | let observed = observed_budget |
| 611 | .lock() |
| 612 | .expect("budget lock") |
| 613 | .expect("first backend must observe a deadline"); |
| 614 | assert!( |
| 615 | observed > Duration::from_millis(1_250), |
| 616 | "dedicated first-attempt budget should exceed the default one-second fair share: {observed:?}" |
| 617 | ); |
| 618 | assert!(observed <= first_attempt_budget); |
| 619 | } |
| 620 | |
| 621 | #[tokio::test] |
| 622 | async fn all_unavailable_returns_typed_error_with_backend_ids_only() { |
| 623 | let private_error = "secret provider response"; |
| 624 | let api = FakeBackend { |
| 625 | id: BackendId::Bocha, |
| 626 | result: Err(ToolError::execution_failed(private_error)), |
| 627 | }; |
| 628 | let scrape = FakeBackend { |
| 629 | id: BackendId::DuckDuckGo, |
| 630 | result: Err(ToolError::execution_failed("different private response")), |
| 631 | }; |
| 632 | let error = run_backend_chain( |
| 633 | &[&api, &scrape], |
| 634 | &query(), |
| 635 | Instant::now() + Duration::from_secs(1), |
| 636 | None, |
| 637 | ) |
| 638 | .await |
| 639 | .expect_err("all-down chain must fail"); |
| 640 | let message = error.to_string(); |
| 641 | |
| 642 | assert!(matches!(error, ToolError::NotAvailable { .. })); |
| 643 | assert!(message.contains("bocha, duckduckgo")); |
| 644 | assert!(!message.contains(private_error)); |
| 645 | assert!(!message.contains("different private response")); |
| 646 | } |
| 647 | |
| 648 | #[tokio::test] |
| 649 | async fn policy_failure_does_not_leak_query_to_fallback() { |
| 650 | struct CountingBackend { |
| 651 | calls: Arc<std::sync::atomic::AtomicUsize>, |
| 652 | } |
| 653 | #[async_trait] |
| 654 | impl SearchBackend for CountingBackend { |
| 655 | fn id(&self) -> BackendId { |
| 656 | BackendId::DuckDuckGo |
| 657 | } |
| 658 | |
| 659 | fn capabilities(&self) -> QueryCapabilities { |
| 660 | QueryCapabilities::count_only() |
| 661 | } |
| 662 | |
| 663 | async fn search( |
| 664 | &self, |
| 665 | _query: &SearchQuery, |
| 666 | _deadline: Instant, |
| 667 | ) -> Result<BackendSearch, ToolError> { |
| 668 | self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); |
| 669 | Err(ToolError::execution_failed("unexpected fallback")) |
| 670 | } |
| 671 | } |
| 672 | |
| 673 | let api = FakeBackend { |
| 674 | id: BackendId::Searxng, |
| 675 | result: Err(ToolError::permission_denied("policy blocked")), |
| 676 | }; |
| 677 | let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); |
| 678 | let scrape = CountingBackend { |
| 679 | calls: Arc::clone(&calls), |
| 680 | }; |
| 681 | let error = run_backend_chain( |
| 682 | &[&api, &scrape], |
| 683 | &query(), |
| 684 | Instant::now() + Duration::from_secs(1), |
| 685 | None, |
| 686 | ) |
| 687 | .await |
| 688 | .expect_err("policy error must fail closed"); |
| 689 | |
| 690 | assert!(matches!(error, ToolError::PermissionDenied { .. })); |
| 691 | assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); |
| 692 | } |
| 693 | |
| 694 | #[tokio::test] |
| 695 | async fn empty_api_falls_back_and_records_no_usable_results() { |
| 696 | let api = FakeBackend { |
| 697 | id: BackendId::Metaso, |
| 698 | result: Ok(Vec::new()), |
| 699 | }; |
| 700 | let scrape = FakeBackend { |
| 701 | id: BackendId::DuckDuckGo, |
| 702 | result: Ok(vec![result()]), |
| 703 | }; |
| 704 | let response = run_backend_chain( |
| 705 | &[&api, &scrape], |
| 706 | &query(), |
| 707 | Instant::now() + Duration::from_secs(1), |
| 708 | None, |
| 709 | ) |
| 710 | .await |
| 711 | .expect("empty API response should fall back"); |
| 712 | |
| 713 | assert_eq!( |
| 714 | response.raw.degraded, |
| 715 | vec![ |
| 716 | DegradedReason::NoUsableResults { |
| 717 | backend: BackendId::Metaso, |
| 718 | }, |
| 719 | DegradedReason::BackendFallback { |
| 720 | from: BackendId::Metaso, |
| 721 | to: BackendId::DuckDuckGo, |
| 722 | }, |
| 723 | ] |
| 724 | ); |
| 725 | } |
| 726 | } |
| 727 |