| 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 | Firecrawl(BackendContext<'a>), |
| 33 | Tavily(BackendContext<'a>), |
| 34 | Bocha(BackendContext<'a>), |
| 35 | Metaso(BackendContext<'a>), |
| 36 | Searxng(BackendContext<'a>), |
| 37 | Baidu(BackendContext<'a>), |
| 38 | Volcengine(BackendContext<'a>), |
| 39 | Sofya(BackendContext<'a>), |
| 40 | Serply(BackendContext<'a>), |
| 41 | } |
| 42 | |
| 43 | #[derive(Clone, Copy)] |
| 44 | struct ProviderNativeSearchBackend<'a> { |
| 45 | context: &'a ToolContext, |
| 46 | } |
| 47 | |
| 48 | impl<'a> ConfiguredSearchBackend<'a> { |
| 49 | #[must_use] |
| 50 | pub(crate) fn from_provider(context: &'a ToolContext, provider: SearchProvider) -> Self { |
| 51 | let backend = BackendContext { |
| 52 | tool_context: context, |
| 53 | }; |
| 54 | match provider { |
| 55 | SearchProvider::Bing => Self::Bing(backend), |
| 56 | SearchProvider::DuckDuckGo => Self::DuckDuckGo(backend), |
| 57 | SearchProvider::Firecrawl => Self::Firecrawl(backend), |
| 58 | SearchProvider::Tavily => Self::Tavily(backend), |
| 59 | SearchProvider::Bocha => Self::Bocha(backend), |
| 60 | SearchProvider::Metaso => Self::Metaso(backend), |
| 61 | SearchProvider::Searxng => Self::Searxng(backend), |
| 62 | SearchProvider::Baidu => Self::Baidu(backend), |
| 63 | SearchProvider::Volcengine => Self::Volcengine(backend), |
| 64 | SearchProvider::Sofya => Self::Sofya(backend), |
| 65 | SearchProvider::Serply => Self::Serply(backend), |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | const fn provider(&self) -> SearchProvider { |
| 70 | match self { |
| 71 | Self::Bing(_) => SearchProvider::Bing, |
| 72 | Self::DuckDuckGo(_) => SearchProvider::DuckDuckGo, |
| 73 | Self::Firecrawl(_) => SearchProvider::Firecrawl, |
| 74 | Self::Tavily(_) => SearchProvider::Tavily, |
| 75 | Self::Bocha(_) => SearchProvider::Bocha, |
| 76 | Self::Metaso(_) => SearchProvider::Metaso, |
| 77 | Self::Searxng(_) => SearchProvider::Searxng, |
| 78 | Self::Baidu(_) => SearchProvider::Baidu, |
| 79 | Self::Volcengine(_) => SearchProvider::Volcengine, |
| 80 | Self::Sofya(_) => SearchProvider::Sofya, |
| 81 | Self::Serply(_) => SearchProvider::Serply, |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | const fn context(&self) -> &BackendContext<'a> { |
| 86 | match self { |
| 87 | Self::Bing(context) |
| 88 | | Self::DuckDuckGo(context) |
| 89 | | Self::Firecrawl(context) |
| 90 | | Self::Tavily(context) |
| 91 | | Self::Bocha(context) |
| 92 | | Self::Metaso(context) |
| 93 | | Self::Searxng(context) |
| 94 | | Self::Baidu(context) |
| 95 | | Self::Volcengine(context) |
| 96 | | Self::Sofya(context) |
| 97 | | Self::Serply(context) => context, |
| 98 | } |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | pub(crate) struct SearchBackendChain<'a> { |
| 103 | backends: Vec<Box<dyn SearchBackend + 'a>>, |
| 104 | } |
| 105 | |
| 106 | #[derive(Debug)] |
| 107 | pub(crate) struct ChainedSearch { |
| 108 | pub(crate) raw: BackendSearch, |
| 109 | pub(crate) capabilities: QueryCapabilities, |
| 110 | } |
| 111 | |
| 112 | impl<'a> SearchBackendChain<'a> { |
| 113 | #[must_use] |
| 114 | pub(crate) fn from_context(context: &'a ToolContext) -> Self { |
| 115 | let selected = context.search_provider; |
| 116 | let mut backends: Vec<Box<dyn SearchBackend + 'a>> = Vec::new(); |
| 117 | if should_prepend_provider_native(context) { |
| 118 | backends.push(Box::new(ProviderNativeSearchBackend { context })); |
| 119 | } |
| 120 | backends.push(Box::new(ConfiguredSearchBackend::from_provider( |
| 121 | context, selected, |
| 122 | ))); |
| 123 | if !matches!(selected, SearchProvider::Bing | SearchProvider::DuckDuckGo) { |
| 124 | backends.push(Box::new(ConfiguredSearchBackend::from_provider( |
| 125 | context, |
| 126 | SearchProvider::DuckDuckGo, |
| 127 | ))); |
| 128 | } |
| 129 | Self { backends } |
| 130 | } |
| 131 | |
| 132 | #[must_use] |
| 133 | pub(crate) fn initial_backend(&self) -> BackendId { |
| 134 | self.backends |
| 135 | .first() |
| 136 | .expect("a search chain always has a configured backend") |
| 137 | .id() |
| 138 | } |
| 139 | |
| 140 | pub(crate) async fn search( |
| 141 | &self, |
| 142 | query: &SearchQuery, |
| 143 | deadline: Instant, |
| 144 | first_attempt_budget: Option<Duration>, |
| 145 | fallback_budget_after_first: Option<Duration>, |
| 146 | ) -> Result<ChainedSearch, ToolError> { |
| 147 | let backends = self |
| 148 | .backends |
| 149 | .iter() |
| 150 | .map(|backend| backend.as_ref()) |
| 151 | .collect::<Vec<_>>(); |
| 152 | run_backend_chain( |
| 153 | &backends, |
| 154 | query, |
| 155 | deadline, |
| 156 | first_attempt_budget, |
| 157 | fallback_budget_after_first, |
| 158 | ) |
| 159 | .await |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | fn should_prepend_provider_native(context: &ToolContext) -> bool { |
| 164 | provider_native_is_available( |
| 165 | context |
| 166 | .route_capabilities |
| 167 | .server_side_web_search |
| 168 | .is_supported(), |
| 169 | context.provider_native_search.is_some(), |
| 170 | ) |
| 171 | } |
| 172 | |
| 173 | const fn provider_native_is_available(capability_supported: bool, client_present: bool) -> bool { |
| 174 | capability_supported && client_present |
| 175 | } |
| 176 | |
| 177 | async fn run_backend_chain( |
| 178 | backends: &[&dyn SearchBackend], |
| 179 | query: &SearchQuery, |
| 180 | mut deadline: Instant, |
| 181 | first_attempt_budget: Option<Duration>, |
| 182 | fallback_budget_after_first: Option<Duration>, |
| 183 | ) -> Result<ChainedSearch, ToolError> { |
| 184 | let mut degraded = Vec::new(); |
| 185 | let mut last_empty = None; |
| 186 | let mut attempted = Vec::new(); |
| 187 | |
| 188 | for (index, backend) in backends.iter().enumerate() { |
| 189 | if index == 1 |
| 190 | && let Some(fallback_budget) = fallback_budget_after_first |
| 191 | { |
| 192 | deadline = Instant::now() + fallback_budget.max(Duration::from_millis(1)); |
| 193 | } |
| 194 | let remaining = deadline.saturating_duration_since(Instant::now()); |
| 195 | if remaining.is_zero() { |
| 196 | break; |
| 197 | } |
| 198 | let backend_id = backend.id(); |
| 199 | if let Some(previous) = attempted.last() { |
| 200 | degraded.push(DegradedReason::BackendFallback { |
| 201 | from: *previous, |
| 202 | to: backend_id, |
| 203 | }); |
| 204 | } |
| 205 | attempted.push(backend_id); |
| 206 | |
| 207 | let attempts_left = u32::try_from(backends.len() - index).unwrap_or(u32::MAX); |
| 208 | let fair_share = remaining / attempts_left; |
| 209 | let attempt_budget = if index == 0 { |
| 210 | first_attempt_budget |
| 211 | .map(|budget| budget.min(remaining)) |
| 212 | .unwrap_or(fair_share) |
| 213 | } else { |
| 214 | fair_share |
| 215 | } |
| 216 | .max(Duration::from_millis(1)); |
| 217 | let attempt_deadline = Instant::now() + attempt_budget; |
| 218 | |
| 219 | let result = tokio::time::timeout(attempt_budget, backend.search(query, attempt_deadline)) |
| 220 | .await |
| 221 | .map_err(|_| ToolError::Timeout { |
| 222 | seconds: u64::try_from(attempt_budget.as_millis()) |
| 223 | .unwrap_or(u64::MAX) |
| 224 | .div_ceil(1_000), |
| 225 | }) |
| 226 | .and_then(std::convert::identity); |
| 227 | |
| 228 | match result { |
| 229 | Ok(mut raw) => { |
| 230 | let capabilities = backend.capabilities(); |
| 231 | crate::tools::web_search::apply_domain_constraints(query, capabilities, &mut raw); |
| 232 | if !raw.results.is_empty() { |
| 233 | degraded.append(&mut raw.degraded); |
| 234 | raw.degraded = degraded; |
| 235 | return Ok(ChainedSearch { raw, capabilities }); |
| 236 | } |
| 237 | degraded.push(DegradedReason::NoUsableResults { |
| 238 | backend: backend_id, |
| 239 | }); |
| 240 | degraded.append(&mut raw.degraded); |
| 241 | last_empty = Some((raw, capabilities)); |
| 242 | } |
| 243 | Err(error) if is_fail_closed(&error) => return Err(error), |
| 244 | Err(error) if backends.len() == 1 => return Err(error), |
| 245 | Err(_) => degraded.push(DegradedReason::BackendUnavailable { |
| 246 | backend: backend_id, |
| 247 | }), |
| 248 | } |
| 249 | } |
| 250 | |
| 251 | if let Some((mut raw, capabilities)) = last_empty { |
| 252 | raw.degraded = degraded; |
| 253 | return Ok(ChainedSearch { raw, capabilities }); |
| 254 | } |
| 255 | |
| 256 | if attempted.is_empty() { |
| 257 | return Err(ToolError::Timeout { seconds: 1 }); |
| 258 | } |
| 259 | |
| 260 | let backend_ids = attempted |
| 261 | .into_iter() |
| 262 | .map(BackendId::as_str) |
| 263 | .collect::<Vec<_>>() |
| 264 | .join(", "); |
| 265 | Err(ToolError::not_available(format!( |
| 266 | "web search backends unavailable: {backend_ids}" |
| 267 | ))) |
| 268 | } |
| 269 | |
| 270 | const fn is_fail_closed(error: &ToolError) -> bool { |
| 271 | matches!( |
| 272 | error, |
| 273 | ToolError::InvalidInput { .. } |
| 274 | | ToolError::MissingField { .. } |
| 275 | | ToolError::PathEscape { .. } |
| 276 | | ToolError::Cancelled { .. } |
| 277 | | ToolError::PermissionDenied { .. } |
| 278 | ) |
| 279 | } |
| 280 | |
| 281 | #[async_trait] |
| 282 | impl SearchBackend for ConfiguredSearchBackend<'_> { |
| 283 | fn id(&self) -> BackendId { |
| 284 | match self.provider() { |
| 285 | SearchProvider::Bing => BackendId::Bing, |
| 286 | SearchProvider::DuckDuckGo => BackendId::DuckDuckGo, |
| 287 | SearchProvider::Firecrawl => BackendId::Firecrawl, |
| 288 | SearchProvider::Tavily => BackendId::Tavily, |
| 289 | SearchProvider::Bocha => BackendId::Bocha, |
| 290 | SearchProvider::Metaso => BackendId::Metaso, |
| 291 | SearchProvider::Searxng => BackendId::Searxng, |
| 292 | SearchProvider::Baidu => BackendId::Baidu, |
| 293 | SearchProvider::Volcengine => BackendId::Volcengine, |
| 294 | SearchProvider::Sofya => BackendId::Sofya, |
| 295 | SearchProvider::Serply => BackendId::Serply, |
| 296 | } |
| 297 | } |
| 298 | |
| 299 | fn capabilities(&self) -> QueryCapabilities { |
| 300 | // All current adapters enforce result count. Other knobs are either |
| 301 | // post-filtered by the shared harness or reported as not honored. |
| 302 | QueryCapabilities::count_only() |
| 303 | } |
| 304 | |
| 305 | async fn search( |
| 306 | &self, |
| 307 | query: &SearchQuery, |
| 308 | deadline: Instant, |
| 309 | ) -> Result<BackendSearch, ToolError> { |
| 310 | crate::tools::web_search::run_backend_search( |
| 311 | self.provider(), |
| 312 | query, |
| 313 | deadline, |
| 314 | self.context().tool_context, |
| 315 | ) |
| 316 | .await |
| 317 | } |
| 318 | } |
| 319 | |
| 320 | #[async_trait] |
| 321 | impl SearchBackend for ProviderNativeSearchBackend<'_> { |
| 322 | fn id(&self) -> BackendId { |
| 323 | BackendId::ProviderNative |
| 324 | } |
| 325 | |
| 326 | fn capabilities(&self) -> QueryCapabilities { |
| 327 | QueryCapabilities { |
| 328 | max_results: QueryCapabilityState::Supported, |
| 329 | recency: QueryCapabilityState::Unsupported, |
| 330 | domains: QueryCapabilityState::Supported, |
| 331 | locale: QueryCapabilityState::Unsupported, |
| 332 | published_date: QueryCapabilityState::Unknown, |
| 333 | } |
| 334 | } |
| 335 | |
| 336 | async fn search( |
| 337 | &self, |
| 338 | query: &SearchQuery, |
| 339 | _deadline: Instant, |
| 340 | ) -> Result<BackendSearch, ToolError> { |
| 341 | if !self |
| 342 | .context |
| 343 | .route_capabilities |
| 344 | .server_side_web_search |
| 345 | .is_supported() |
| 346 | { |
| 347 | return Err(ToolError::not_available( |
| 348 | "active route does not report provider-native web search", |
| 349 | )); |
| 350 | } |
| 351 | let client = self |
| 352 | .context |
| 353 | .provider_native_search |
| 354 | .as_ref() |
| 355 | .ok_or_else(|| ToolError::not_available("provider-native search client unavailable"))?; |
| 356 | // Moonshot/Kimi, Z.AI, MiMo, and the Responses-dialect routes cannot |
| 357 | // express domain filters in their native wire contracts. Declining |
| 358 | // here must not fail the whole search: report this backend unavailable |
| 359 | // so the chain falls back to the configured provider or DuckDuckGo, |
| 360 | // which honor domains natively or through post-filtering. |
| 361 | let domain_limit = client.maximum_domain_count(); |
| 362 | if !query.domains.is_empty() && domain_limit == Some(0) { |
| 363 | return Err(ToolError::not_available(format!( |
| 364 | "{} native web search cannot honor domain filters", |
| 365 | client.provider().as_str() |
| 366 | ))); |
| 367 | } |
| 368 | if let Some(maximum) = domain_limit |
| 369 | && query.domains.len() > maximum |
| 370 | { |
| 371 | return Err(ToolError::invalid_input(format!( |
| 372 | "{} native web search accepts at most {maximum} domains", |
| 373 | client.provider().as_str() |
| 374 | ))); |
| 375 | } |
| 376 | let host = client.host().ok_or_else(|| { |
| 377 | ToolError::execution_failed("provider-native search endpoint has no valid host") |
| 378 | })?; |
| 379 | crate::tools::web_search::check_policy( |
| 380 | self.context.network_policy.as_ref(), |
| 381 | host.as_str(), |
| 382 | )?; |
| 383 | let response = client |
| 384 | .search(&ProviderNativeSearchRequest { |
| 385 | query: query.query.clone(), |
| 386 | max_results: query.max_results, |
| 387 | domains: query.domains.clone(), |
| 388 | }) |
| 389 | .await |
| 390 | .map_err(|error| { |
| 391 | ToolError::execution_failed(format!( |
| 392 | "{} provider-native web search failed: {error}", |
| 393 | client.provider().as_str() |
| 394 | )) |
| 395 | })?; |
| 396 | let results = response |
| 397 | .citations |
| 398 | .into_iter() |
| 399 | .enumerate() |
| 400 | .map(|(index, citation)| { |
| 401 | SearchResult::new( |
| 402 | index + 1, |
| 403 | citation.title, |
| 404 | citation.url, |
| 405 | citation.snippet, |
| 406 | citation.published, |
| 407 | ) |
| 408 | }) |
| 409 | .collect(); |
| 410 | Ok(BackendSearch { |
| 411 | backend: BackendId::ProviderNative, |
| 412 | source: format!( |
| 413 | "provider-native/{}/{}", |
| 414 | client.provider().as_str(), |
| 415 | client.model() |
| 416 | ), |
| 417 | backend_detail: Some(host), |
| 418 | results, |
| 419 | degraded: Vec::new(), |
| 420 | note: response.answer, |
| 421 | }) |
| 422 | } |
| 423 | } |
| 424 | |
| 425 | #[cfg(test)] |
| 426 | mod tests { |
| 427 | use std::sync::{Arc, Mutex}; |
| 428 | |
| 429 | use super::*; |
| 430 | |
| 431 | struct FakeBackend { |
| 432 | id: BackendId, |
| 433 | result: Result<Vec<super::super::contract::SearchResult>, ToolError>, |
| 434 | } |
| 435 | |
| 436 | struct DeadlineBackend { |
| 437 | id: BackendId, |
| 438 | observed_budget: Arc<Mutex<Option<Duration>>>, |
| 439 | delay: Duration, |
| 440 | } |
| 441 | |
| 442 | #[async_trait] |
| 443 | impl SearchBackend for FakeBackend { |
| 444 | fn id(&self) -> BackendId { |
| 445 | self.id |
| 446 | } |
| 447 | |
| 448 | fn capabilities(&self) -> QueryCapabilities { |
| 449 | QueryCapabilities::count_only() |
| 450 | } |
| 451 | |
| 452 | async fn search( |
| 453 | &self, |
| 454 | _query: &SearchQuery, |
| 455 | _deadline: Instant, |
| 456 | ) -> Result<BackendSearch, ToolError> { |
| 457 | Ok(BackendSearch { |
| 458 | backend: self.id, |
| 459 | source: self.id.as_str().to_string(), |
| 460 | backend_detail: None, |
| 461 | results: self.result.clone()?, |
| 462 | degraded: Vec::new(), |
| 463 | note: None, |
| 464 | }) |
| 465 | } |
| 466 | } |
| 467 | |
| 468 | #[async_trait] |
| 469 | impl SearchBackend for DeadlineBackend { |
| 470 | fn id(&self) -> BackendId { |
| 471 | self.id |
| 472 | } |
| 473 | |
| 474 | fn capabilities(&self) -> QueryCapabilities { |
| 475 | QueryCapabilities::count_only() |
| 476 | } |
| 477 | |
| 478 | async fn search( |
| 479 | &self, |
| 480 | _query: &SearchQuery, |
| 481 | deadline: Instant, |
| 482 | ) -> Result<BackendSearch, ToolError> { |
| 483 | *self.observed_budget.lock().expect("budget lock") = |
| 484 | Some(deadline.saturating_duration_since(Instant::now())); |
| 485 | tokio::time::sleep(self.delay).await; |
| 486 | Ok(BackendSearch { |
| 487 | backend: self.id, |
| 488 | source: self.id.as_str().to_string(), |
| 489 | backend_detail: None, |
| 490 | results: vec![result()], |
| 491 | degraded: Vec::new(), |
| 492 | note: None, |
| 493 | }) |
| 494 | } |
| 495 | } |
| 496 | |
| 497 | fn query() -> SearchQuery { |
| 498 | SearchQuery::new("bounded chain".to_string(), 5, None, Vec::new(), None) |
| 499 | } |
| 500 | |
| 501 | fn result() -> super::super::contract::SearchResult { |
| 502 | super::super::contract::SearchResult::new( |
| 503 | 1, |
| 504 | "Fallback result".to_string(), |
| 505 | "https://example.com/result".to_string(), |
| 506 | None, |
| 507 | None, |
| 508 | ) |
| 509 | } |
| 510 | |
| 511 | #[test] |
| 512 | fn every_configured_provider_maps_to_one_explicit_backend_adapter() { |
| 513 | let cases = [ |
| 514 | (SearchProvider::Bing, BackendId::Bing), |
| 515 | (SearchProvider::DuckDuckGo, BackendId::DuckDuckGo), |
| 516 | (SearchProvider::Firecrawl, BackendId::Firecrawl), |
| 517 | (SearchProvider::Tavily, BackendId::Tavily), |
| 518 | (SearchProvider::Bocha, BackendId::Bocha), |
| 519 | (SearchProvider::Metaso, BackendId::Metaso), |
| 520 | (SearchProvider::Searxng, BackendId::Searxng), |
| 521 | (SearchProvider::Baidu, BackendId::Baidu), |
| 522 | (SearchProvider::Volcengine, BackendId::Volcengine), |
| 523 | (SearchProvider::Sofya, BackendId::Sofya), |
| 524 | (SearchProvider::Serply, BackendId::Serply), |
| 525 | ]; |
| 526 | |
| 527 | for (provider, expected) in cases { |
| 528 | let mut context = ToolContext::new(std::path::PathBuf::from(".")); |
| 529 | context.search_provider = provider; |
| 530 | let backend = ConfiguredSearchBackend::from_provider(&context, provider); |
| 531 | assert_eq!(backend.id(), expected); |
| 532 | assert_eq!( |
| 533 | backend.capabilities().max_results, |
| 534 | super::super::contract::CapabilityState::Supported |
| 535 | ); |
| 536 | } |
| 537 | } |
| 538 | |
| 539 | #[test] |
| 540 | fn provider_native_is_fail_closed_without_both_fact_and_client() { |
| 541 | assert!(!provider_native_is_available(false, false)); |
| 542 | assert!(!provider_native_is_available(true, false)); |
| 543 | assert!(!provider_native_is_available(false, true)); |
| 544 | assert!(provider_native_is_available(true, true)); |
| 545 | } |
| 546 | |
| 547 | #[tokio::test] |
| 548 | async fn unavailable_api_falls_back_with_explicit_receipts() { |
| 549 | let api = FakeBackend { |
| 550 | id: BackendId::Tavily, |
| 551 | result: Err(ToolError::execution_failed( |
| 552 | "provider detail must stay private", |
| 553 | )), |
| 554 | }; |
| 555 | let scrape = FakeBackend { |
| 556 | id: BackendId::DuckDuckGo, |
| 557 | result: Ok(vec![result()]), |
| 558 | }; |
| 559 | let response = run_backend_chain( |
| 560 | &[&api, &scrape], |
| 561 | &query(), |
| 562 | Instant::now() + Duration::from_secs(1), |
| 563 | None, |
| 564 | None, |
| 565 | ) |
| 566 | .await |
| 567 | .expect("fallback should succeed"); |
| 568 | |
| 569 | assert_eq!(response.raw.backend, BackendId::DuckDuckGo); |
| 570 | assert_eq!( |
| 571 | response.raw.degraded, |
| 572 | vec![ |
| 573 | DegradedReason::BackendUnavailable { |
| 574 | backend: BackendId::Tavily, |
| 575 | }, |
| 576 | DegradedReason::BackendFallback { |
| 577 | from: BackendId::Tavily, |
| 578 | to: BackendId::DuckDuckGo, |
| 579 | }, |
| 580 | ] |
| 581 | ); |
| 582 | } |
| 583 | |
| 584 | #[tokio::test] |
| 585 | async fn provider_native_to_api_to_scrape_records_every_transition() { |
| 586 | let native = FakeBackend { |
| 587 | id: BackendId::ProviderNative, |
| 588 | result: Err(ToolError::execution_failed("native unavailable")), |
| 589 | }; |
| 590 | let api = FakeBackend { |
| 591 | id: BackendId::Tavily, |
| 592 | result: Err(ToolError::execution_failed("API unavailable")), |
| 593 | }; |
| 594 | let scrape = FakeBackend { |
| 595 | id: BackendId::DuckDuckGo, |
| 596 | result: Ok(vec![result()]), |
| 597 | }; |
| 598 | |
| 599 | let response = run_backend_chain( |
| 600 | &[&native, &api, &scrape], |
| 601 | &query(), |
| 602 | Instant::now() + Duration::from_secs(1), |
| 603 | None, |
| 604 | None, |
| 605 | ) |
| 606 | .await |
| 607 | .expect("final scrape fallback should succeed"); |
| 608 | |
| 609 | assert_eq!(response.raw.backend, BackendId::DuckDuckGo); |
| 610 | assert_eq!( |
| 611 | response.raw.degraded, |
| 612 | vec![ |
| 613 | DegradedReason::BackendUnavailable { |
| 614 | backend: BackendId::ProviderNative, |
| 615 | }, |
| 616 | DegradedReason::BackendFallback { |
| 617 | from: BackendId::ProviderNative, |
| 618 | to: BackendId::Tavily, |
| 619 | }, |
| 620 | DegradedReason::BackendUnavailable { |
| 621 | backend: BackendId::Tavily, |
| 622 | }, |
| 623 | DegradedReason::BackendFallback { |
| 624 | from: BackendId::Tavily, |
| 625 | to: BackendId::DuckDuckGo, |
| 626 | }, |
| 627 | ] |
| 628 | ); |
| 629 | } |
| 630 | |
| 631 | #[tokio::test] |
| 632 | async fn zero_domain_native_providers_decline_without_failing_the_chain() { |
| 633 | use crate::config::{Config, ProviderConfig, ProvidersConfig}; |
| 634 | |
| 635 | let moonshot_config = Config { |
| 636 | provider: Some("moonshot".to_string()), |
| 637 | providers: Some(ProvidersConfig { |
| 638 | moonshot: ProviderConfig { |
| 639 | api_key: Some("moonshot-test-key".to_string()), |
| 640 | base_url: Some("https://api.moonshot.ai/v1".to_string()), |
| 641 | model: Some("kimi-k3".to_string()), |
| 642 | ..ProviderConfig::default() |
| 643 | }, |
| 644 | ..ProvidersConfig::default() |
| 645 | }), |
| 646 | ..Config::default() |
| 647 | }; |
| 648 | let tmp = tempfile::tempdir().expect("tempdir"); |
| 649 | let mut context = ToolContext::new(tmp.path().to_path_buf()); |
| 650 | context.route_capabilities.server_side_web_search = |
| 651 | codewhale_config::route::CapabilityState::Supported; |
| 652 | context.provider_native_search = Some( |
| 653 | crate::client::ProviderNativeSearchClient::new( |
| 654 | crate::client::CodewhaleClient::new(&moonshot_config) |
| 655 | .expect("test Moonshot client"), |
| 656 | ) |
| 657 | .expect("Moonshot native adapter"), |
| 658 | ); |
| 659 | let backend = ProviderNativeSearchBackend { context: &context }; |
| 660 | |
| 661 | let domain_query = SearchQuery::new( |
| 662 | "bounded chain".to_string(), |
| 663 | 5, |
| 664 | None, |
| 665 | vec!["example.com".to_string()], |
| 666 | None, |
| 667 | ); |
| 668 | let error = backend |
| 669 | .search(&domain_query, Instant::now() + Duration::from_secs(1)) |
| 670 | .await |
| 671 | .expect_err("Moonshot native search must decline domain-filtered queries"); |
| 672 | assert!( |
| 673 | matches!(error, ToolError::NotAvailable { .. }), |
| 674 | "declining must stay fallback-shaped, not fail-closed: {error:?}" |
| 675 | ); |
| 676 | |
| 677 | let xai_config = Config { |
| 678 | provider: Some("xai".to_string()), |
| 679 | providers: Some(ProvidersConfig { |
| 680 | xai: ProviderConfig { |
| 681 | api_key: Some("xai-test-key".to_string()), |
| 682 | base_url: Some("https://api.x.ai/v1".to_string()), |
| 683 | model: Some("grok-4.5".to_string()), |
| 684 | ..ProviderConfig::default() |
| 685 | }, |
| 686 | ..ProvidersConfig::default() |
| 687 | }), |
| 688 | ..Config::default() |
| 689 | }; |
| 690 | let mut xai_context = ToolContext::new(tmp.path().to_path_buf()); |
| 691 | xai_context.route_capabilities.server_side_web_search = |
| 692 | codewhale_config::route::CapabilityState::Supported; |
| 693 | xai_context.provider_native_search = Some( |
| 694 | crate::client::ProviderNativeSearchClient::new( |
| 695 | crate::client::CodewhaleClient::new(&xai_config).expect("test xAI client"), |
| 696 | ) |
| 697 | .expect("xAI native adapter"), |
| 698 | ); |
| 699 | let oversized_domain_query = SearchQuery::new( |
| 700 | "bounded chain".to_string(), |
| 701 | 5, |
| 702 | None, |
| 703 | [ |
| 704 | "a.example", |
| 705 | "b.example", |
| 706 | "c.example", |
| 707 | "d.example", |
| 708 | "e.example", |
| 709 | "f.example", |
| 710 | ] |
| 711 | .iter() |
| 712 | .map(|domain| domain.to_string()) |
| 713 | .collect(), |
| 714 | None, |
| 715 | ); |
| 716 | let error = ProviderNativeSearchBackend { |
| 717 | context: &xai_context, |
| 718 | } |
| 719 | .search( |
| 720 | &oversized_domain_query, |
| 721 | Instant::now() + Duration::from_secs(1), |
| 722 | ) |
| 723 | .await |
| 724 | .expect_err("too many domains stays a typed user error"); |
| 725 | assert!( |
| 726 | matches!(error, ToolError::InvalidInput { .. }), |
| 727 | "over the provider limit must stay fail-closed: {error:?}" |
| 728 | ); |
| 729 | } |
| 730 | |
| 731 | #[tokio::test] |
| 732 | async fn first_attempt_budget_overrides_the_default_fair_share() { |
| 733 | let observed_budget = Arc::new(Mutex::new(None)); |
| 734 | let volcengine = DeadlineBackend { |
| 735 | id: BackendId::Volcengine, |
| 736 | observed_budget: Arc::clone(&observed_budget), |
| 737 | delay: Duration::ZERO, |
| 738 | }; |
| 739 | let fallback = FakeBackend { |
| 740 | id: BackendId::DuckDuckGo, |
| 741 | result: Ok(vec![result()]), |
| 742 | }; |
| 743 | let first_attempt_budget = Duration::from_millis(1_500); |
| 744 | let response = run_backend_chain( |
| 745 | &[&volcengine, &fallback], |
| 746 | &query(), |
| 747 | Instant::now() + Duration::from_secs(2), |
| 748 | Some(first_attempt_budget), |
| 749 | None, |
| 750 | ) |
| 751 | .await |
| 752 | .expect("the first backend should complete inside its dedicated budget"); |
| 753 | |
| 754 | assert_eq!(response.raw.backend, BackendId::Volcengine); |
| 755 | let observed = observed_budget |
| 756 | .lock() |
| 757 | .expect("budget lock") |
| 758 | .expect("first backend must observe a deadline"); |
| 759 | assert!( |
| 760 | observed > Duration::from_millis(1_250), |
| 761 | "dedicated first-attempt budget should exceed the default one-second fair share: {observed:?}" |
| 762 | ); |
| 763 | assert!(observed <= first_attempt_budget); |
| 764 | } |
| 765 | |
| 766 | #[tokio::test] |
| 767 | async fn provider_native_unused_budget_does_not_extend_fallback_deadline() { |
| 768 | let native = FakeBackend { |
| 769 | id: BackendId::ProviderNative, |
| 770 | result: Err(ToolError::execution_failed("native unavailable")), |
| 771 | }; |
| 772 | let observed_budget = Arc::new(Mutex::new(None)); |
| 773 | let fallback = DeadlineBackend { |
| 774 | id: BackendId::DuckDuckGo, |
| 775 | observed_budget: Arc::clone(&observed_budget), |
| 776 | delay: Duration::from_millis(200), |
| 777 | }; |
| 778 | let fallback_budget = Duration::from_millis(30); |
| 779 | let error = run_backend_chain( |
| 780 | &[&native, &fallback], |
| 781 | &query(), |
| 782 | Instant::now() + Duration::from_millis(500), |
| 783 | Some(Duration::from_millis(500)), |
| 784 | Some(fallback_budget), |
| 785 | ) |
| 786 | .await |
| 787 | .expect_err("blocking fallback must stop at its own budget"); |
| 788 | |
| 789 | assert!(matches!(error, ToolError::NotAvailable { .. })); |
| 790 | let observed = observed_budget |
| 791 | .lock() |
| 792 | .expect("budget lock") |
| 793 | .expect("fallback must observe a deadline"); |
| 794 | assert!(observed <= fallback_budget); |
| 795 | } |
| 796 | |
| 797 | #[tokio::test] |
| 798 | async fn all_unavailable_returns_typed_error_with_backend_ids_only() { |
| 799 | let private_error = "secret provider response"; |
| 800 | let api = FakeBackend { |
| 801 | id: BackendId::Bocha, |
| 802 | result: Err(ToolError::execution_failed(private_error)), |
| 803 | }; |
| 804 | let scrape = FakeBackend { |
| 805 | id: BackendId::DuckDuckGo, |
| 806 | result: Err(ToolError::execution_failed("different private response")), |
| 807 | }; |
| 808 | let error = run_backend_chain( |
| 809 | &[&api, &scrape], |
| 810 | &query(), |
| 811 | Instant::now() + Duration::from_secs(1), |
| 812 | None, |
| 813 | None, |
| 814 | ) |
| 815 | .await |
| 816 | .expect_err("all-down chain must fail"); |
| 817 | let message = error.to_string(); |
| 818 | |
| 819 | assert!(matches!(error, ToolError::NotAvailable { .. })); |
| 820 | assert!(message.contains("bocha, duckduckgo")); |
| 821 | assert!(!message.contains(private_error)); |
| 822 | assert!(!message.contains("different private response")); |
| 823 | } |
| 824 | |
| 825 | #[tokio::test] |
| 826 | async fn policy_failure_does_not_leak_query_to_fallback() { |
| 827 | struct CountingBackend { |
| 828 | calls: Arc<std::sync::atomic::AtomicUsize>, |
| 829 | } |
| 830 | #[async_trait] |
| 831 | impl SearchBackend for CountingBackend { |
| 832 | fn id(&self) -> BackendId { |
| 833 | BackendId::DuckDuckGo |
| 834 | } |
| 835 | |
| 836 | fn capabilities(&self) -> QueryCapabilities { |
| 837 | QueryCapabilities::count_only() |
| 838 | } |
| 839 | |
| 840 | async fn search( |
| 841 | &self, |
| 842 | _query: &SearchQuery, |
| 843 | _deadline: Instant, |
| 844 | ) -> Result<BackendSearch, ToolError> { |
| 845 | self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); |
| 846 | Err(ToolError::execution_failed("unexpected fallback")) |
| 847 | } |
| 848 | } |
| 849 | |
| 850 | let api = FakeBackend { |
| 851 | id: BackendId::Searxng, |
| 852 | result: Err(ToolError::permission_denied("policy blocked")), |
| 853 | }; |
| 854 | let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); |
| 855 | let scrape = CountingBackend { |
| 856 | calls: Arc::clone(&calls), |
| 857 | }; |
| 858 | let error = run_backend_chain( |
| 859 | &[&api, &scrape], |
| 860 | &query(), |
| 861 | Instant::now() + Duration::from_secs(1), |
| 862 | None, |
| 863 | None, |
| 864 | ) |
| 865 | .await |
| 866 | .expect_err("policy error must fail closed"); |
| 867 | |
| 868 | assert!(matches!(error, ToolError::PermissionDenied { .. })); |
| 869 | assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); |
| 870 | } |
| 871 | |
| 872 | #[tokio::test] |
| 873 | async fn empty_api_falls_back_and_records_no_usable_results() { |
| 874 | let api = FakeBackend { |
| 875 | id: BackendId::Metaso, |
| 876 | result: Ok(Vec::new()), |
| 877 | }; |
| 878 | let scrape = FakeBackend { |
| 879 | id: BackendId::DuckDuckGo, |
| 880 | result: Ok(vec![result()]), |
| 881 | }; |
| 882 | let response = run_backend_chain( |
| 883 | &[&api, &scrape], |
| 884 | &query(), |
| 885 | Instant::now() + Duration::from_secs(1), |
| 886 | None, |
| 887 | None, |
| 888 | ) |
| 889 | .await |
| 890 | .expect("empty API response should fall back"); |
| 891 | |
| 892 | assert_eq!( |
| 893 | response.raw.degraded, |
| 894 | vec![ |
| 895 | DegradedReason::NoUsableResults { |
| 896 | backend: BackendId::Metaso, |
| 897 | }, |
| 898 | DegradedReason::BackendFallback { |
| 899 | from: BackendId::Metaso, |
| 900 | to: BackendId::DuckDuckGo, |
| 901 | }, |
| 902 | ] |
| 903 | ); |
| 904 | } |
| 905 | |
| 906 | #[tokio::test] |
| 907 | async fn domain_filtered_results_fall_back_before_chain_success() { |
| 908 | let native = FakeBackend { |
| 909 | id: BackendId::ProviderNative, |
| 910 | result: Ok(vec![SearchResult::new( |
| 911 | 1, |
| 912 | "Outside source".to_string(), |
| 913 | "https://outside.test/result".to_string(), |
| 914 | None, |
| 915 | None, |
| 916 | )]), |
| 917 | }; |
| 918 | let configured = FakeBackend { |
| 919 | id: BackendId::Searxng, |
| 920 | result: Ok(vec![SearchResult::new( |
| 921 | 1, |
| 922 | "Matching source".to_string(), |
| 923 | "https://docs.rs/example/latest/example/".to_string(), |
| 924 | None, |
| 925 | None, |
| 926 | )]), |
| 927 | }; |
| 928 | let constrained = SearchQuery::new( |
| 929 | "example docs".to_string(), |
| 930 | 5, |
| 931 | None, |
| 932 | vec!["docs.rs".to_string()], |
| 933 | None, |
| 934 | ); |
| 935 | |
| 936 | let response = run_backend_chain( |
| 937 | &[&native, &configured], |
| 938 | &constrained, |
| 939 | Instant::now() + Duration::from_secs(1), |
| 940 | None, |
| 941 | None, |
| 942 | ) |
| 943 | .await |
| 944 | .expect("configured backend should satisfy the domain constraint"); |
| 945 | |
| 946 | assert_eq!(response.raw.backend, BackendId::Searxng); |
| 947 | assert_eq!(response.raw.results.len(), 1); |
| 948 | assert!(response.raw.degraded.iter().any(|reason| matches!( |
| 949 | reason, |
| 950 | DegradedReason::NoUsableResults { |
| 951 | backend: BackendId::ProviderNative |
| 952 | } |
| 953 | ))); |
| 954 | assert!(response.raw.degraded.iter().any(|reason| matches!( |
| 955 | reason, |
| 956 | DegradedReason::BackendFallback { |
| 957 | from: BackendId::ProviderNative, |
| 958 | to: BackendId::Searxng |
| 959 | } |
| 960 | ))); |
| 961 | } |
| 962 | } |
| 963 |