Coverage for src/local_deep_research/web_search_engines/engine_availability.py: 97%
250 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""Research-scoped search engine availability, independent of agent frameworks.
3Empty results are not failures. Engines report an explicit ``SearchFailure``
4when a request cannot work; callers decide when to apply it (after retries and
5fallbacks have finished). One tracker is shared by a research run's workers,
6never by unrelated users or runs. No settings or credentials are persisted.
7"""
9from __future__ import annotations
11import math
12import re
13import threading
14import time
15from contextlib import contextmanager
16from dataclasses import dataclass
17from datetime import UTC, datetime
18from email.utils import parsedate_to_datetime
19from typing import Callable
21import requests
22from urllib3.exceptions import ResponseError
24from ..security.secure_logging import logger
25from .rate_limiting.exceptions import RateLimitError
27# Total time one research run (one availability tracker) spends waiting for
28# cooldowns. Once it is used up, cooling engines are disabled for the rest
29# of the run instead of being waited for again, so a persistently failing
30# sole engine cannot keep a run waiting indefinitely. Users can adjust this
31# per research run with langgraph_agent.cooldown_wait_budget_seconds.
32COOLDOWN_WAIT_BUDGET_SECONDS = 3600
34# Longest provider Retry-After kept as a temporary cooldown; a larger wait
35# budget raises this limit. A longer delay is not one this research can plan
36# around (an HTTP date years away parses to centuries of seconds), so the
37# engine is disabled for the rest of the run with that reason, instead of
38# being shown as "temporary" with an absurd retry time.
39MAX_TEMPORARY_COOLDOWN_SECONDS = 24 * 60 * 60
41# HTTP 403 is often a temporary bot/WAF block rather than bad credentials;
42# unless the reply says it is about authentication, cool down this long.
43FORBIDDEN_COOLDOWN_SECONDS = 300
45# Words in a 403 reply that mark it as an authentication or billing refusal.
46_AUTH_MARKERS = re.compile(
47 r"api[ _-]?key|access[ _-]?token|invalid[ _-]?token|unauthori[sz]ed|"
48 r"authenticat|credential|subscription|billing|payment",
49 re.IGNORECASE,
50)
53@dataclass(frozen=True)
54class SearchFailure:
55 """Safe, actionable reason and cooldown; None means the rest of this run."""
57 reason: str
58 cooldown_seconds: float | None = 60
61def retry_after_seconds(value, *, now: float | None = None) -> float | None:
62 """Parse HTTP Retry-After (seconds or HTTP date), ignoring invalid values."""
63 if not isinstance(value, (str, int, float)) or isinstance(value, bool):
64 return None
65 try:
66 seconds = float(value)
67 except OverflowError:
68 return None
69 except (ValueError, TypeError):
70 if not isinstance(value, str): 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true
71 return None
72 try:
73 seconds = parsedate_to_datetime(value).timestamp() - (
74 time.time() if now is None else now
75 )
76 except (ValueError, TypeError, OverflowError):
77 return None
78 if not math.isfinite(seconds) or seconds <= 0:
79 return None
80 # Provider headers are untrusted; an absurd finite delta must not break
81 # timestamp formatting in every subsequent progress/model update.
82 try:
83 datetime.fromtimestamp(
84 (time.time() if now is None else now) + seconds, UTC
85 )
86 except (OverflowError, OSError, ValueError):
87 return None
88 return seconds
91def retry_after_from_error(error: object) -> float | None:
92 """Preserve Retry-After when providers translate their HTTP exceptions."""
93 delay = retry_after_seconds(getattr(error, "retry_after", None))
94 if delay is None:
95 response = getattr(error, "response", None)
96 headers = getattr(response, "headers", None) or {}
97 delay = retry_after_seconds(headers.get("Retry-After"))
98 return delay
101# HTTP statuses caused by the request's query (malformed syntax, too long,
102# unprocessable), not by the provider being down or refusing this client.
103QUERY_ERROR_STATUSES = frozenset({400, 414, 422})
106def http_status(error: BaseException) -> int | None:
107 """The HTTP status an exception carries, if any."""
108 status = getattr(getattr(error, "response", None), "status_code", None)
109 if not isinstance(status, int):
110 status = getattr(error, "status_code", None)
111 return status if isinstance(status, int) else None
114def is_query_error(error: BaseException) -> bool:
115 """True when the failure is specific to this query, not to the engine.
117 Such an error must not cool the engine down for every other search in
118 the run; a different query can still succeed right away.
119 """
120 return http_status(error) in QUERY_ERROR_STATUSES
123def failure_from_exception(
124 error: BaseException, *, bare_403_is_bad_key: bool = False
125) -> SearchFailure | None:
126 """Classify explicit failures without putting provider text/secrets in UI.
128 ``bare_403_is_bad_key``: the failed request carried an API key to a
129 provider known to reject an invalid key with a plain HTTP 403 (no
130 ``WWW-Authenticate``, nothing telling in the body), so any 403 from it
131 is treated as denied credentials. Other 403s stay a temporary block
132 unless the reply itself says it is about credentials.
133 """
134 status = http_status(error)
135 if isinstance(error, RateLimitError) or status == 429:
136 return SearchFailure(
137 "Rate limit reached", retry_after_from_error(error) or 300
138 )
139 if status in (401, 402) or (
140 status == 403 and (bare_403_is_bad_key or _is_auth_refusal(error))
141 ):
142 return SearchFailure(
143 f"Provider denied access (HTTP {status}); check credentials, permissions or quota",
144 None,
145 )
146 if status == 403:
147 return SearchFailure(
148 "Provider refused the request (HTTP 403); this is often a "
149 "temporary block",
150 FORBIDDEN_COOLDOWN_SECONDS,
151 )
152 if isinstance(status, int) and status >= 500:
153 return SearchFailure(
154 f"Provider unavailable (HTTP {status})",
155 (retry_after_from_error(error) if status == 503 else None) or 60,
156 )
157 if _retries_exhausted_on_429(error):
158 # urllib3 already retried (honoring Retry-After) and discarded the
159 # response, so no provider delay is left to preserve.
160 return SearchFailure("Rate limit reached", 300)
161 if isinstance(error, (requests.Timeout, TimeoutError)):
162 return SearchFailure("Search request timed out")
163 if isinstance(error, (requests.ConnectionError, ConnectionError)):
164 return SearchFailure("Could not connect to the search provider")
165 if isinstance(
166 error,
167 (
168 requests.exceptions.ChunkedEncodingError,
169 requests.exceptions.ContentDecodingError,
170 ),
171 ):
172 return SearchFailure("Could not read the provider response")
173 if isinstance(error, ValueError) and "security validation" in str(error):
174 # The SSRF validator also refuses a host it cannot resolve, and
175 # its ValueError does not say which. A transient DNS failure must
176 # not take an engine out for the whole run, so use a finite
177 # cooldown. A retry is not free: validation resolves the host, and
178 # the redirect-hop form of this error is raised only after the
179 # provider request was sent. Engines that can identify a policy
180 # block (SearXNG's unapproved private instance URL) report a
181 # run-long SearchFailure themselves.
182 return SearchFailure(
183 "Request URL blocked by security validation or its host could "
184 "not be resolved; check server configuration and DNS",
185 300,
186 )
187 return None
190def _is_auth_refusal(error: BaseException) -> bool:
191 """Whether a 403 reply says it is about credentials, not a block.
193 Reads only the response (never the exception text, whose request URL
194 can carry ``api_key=``) and never surfaces what it reads.
195 """
196 response = getattr(error, "response", None)
197 headers = getattr(response, "headers", None) or {}
198 try:
199 if headers.get("WWW-Authenticate"):
200 return True
201 text = getattr(response, "text", "")
202 except Exception:
203 return False
204 return isinstance(text, str) and bool(_AUTH_MARKERS.search(text[:4000]))
207def _retries_exhausted_on_429(error: BaseException) -> bool:
208 """A mounted ``Retry(status_forcelist=[429, ...])`` gave up on HTTP 429.
210 requests raises ``RetryError(MaxRetryError(reason=ResponseError(...)))``
211 without a response; only the reason text names the status. Exhausted
212 retries on other statuses stay generic transient failures.
213 """
214 if not isinstance(error, requests.exceptions.RetryError):
215 return False
216 cause = error.args[0] if error.args else None
217 reason = getattr(cause, "reason", None)
218 return isinstance(reason, ResponseError) and str(
219 reason
220 ) == ResponseError.SPECIFIC_ERROR.format(status_code=429)
223@dataclass(frozen=True)
224class UnrecordedFailure:
225 """A failure that could not be attributed to an engine; not tracked."""
227 failure: SearchFailure
229 def describe(self, name: str | None = None) -> str:
230 return f"{name or 'Search'} failed: {self.failure.reason}"
233@dataclass(frozen=True)
234class EngineDisablement:
235 engine: str
236 failure: SearchFailure
237 disabled_at: float
238 retry_at: float | None
239 deadline: float | None
241 def describe(self, name: str | None = None) -> str:
242 start = datetime.fromtimestamp(self.disabled_at, UTC).strftime(
243 "%Y-%m-%d %H:%M:%S UTC"
244 )
245 if self.retry_at is None:
246 timing = "for the rest of this research; fix configuration before the next run"
247 else:
248 end = datetime.fromtimestamp(self.retry_at, UTC).strftime(
249 "%Y-%m-%d %H:%M:%S UTC"
250 )
251 duration = math.ceil(self.retry_at - self.disabled_at)
252 timing = f"for {duration}s; retry available at {end}"
253 return (
254 f"{name or self.engine} disabled since {start} {timing}: "
255 f"{self.failure.reason}"
256 )
259class SearchEngineAvailability:
260 """Thread-safe circuit breaker with lazy cooldown expiry and change events.
262 No lock is held while a search runs: parallel searches, including the
263 first calls to an engine and retries after a cooldown, run concurrently
264 and record their outcome under the tracker's lock. The first failure
265 wins; later failures of the same outage return the existing state. A
266 failure blocks new calls; work already in flight may still return
267 usable results. Expiry permits a real retry, not a claim that the
268 provider has recovered.
269 """
271 def __init__(
272 self,
273 on_change: Callable | None = None,
274 *,
275 clock: Callable[[], float] = time.monotonic,
276 wall_clock: Callable[[], float] = time.time,
277 check_cancelled: Callable[[], None] | None = None,
278 wait_budget_seconds: float = COOLDOWN_WAIT_BUDGET_SECONDS,
279 max_consecutive_cooldowns: int | None = None,
280 ):
281 self._clock = clock
282 self._wall_clock = wall_clock
283 self._on_change = on_change
284 # Raises (e.g. ResearchTerminatedException) when the run is
285 # cancelled; cooldown waits call it before every sleep. In a research
286 # run the call that sees the cancel first also clears the run's
287 # single-use termination flag while cleaning up; see note_cancelled.
288 self.check_cancelled = check_cancelled
289 # Set once any thread sharing this tracker has seen the run cancelled.
290 self._cancelled = threading.Event()
291 self._lock = threading.RLock()
292 self._disabled: dict[str, EngineDisablement] = {}
293 self._metrics_names: dict[str, str] = {}
294 self._wait_budget = wait_budget_seconds
295 self._max_consecutive_cooldowns = max_consecutive_cooldowns
296 self._consecutive_cooldowns: dict[str, int] = {}
297 # Wall time with at least one waiter, so parallel workers waiting
298 # for the same cooldown spend the budget once, not once each.
299 self._waited = 0.0
300 self._waiters = 0
301 self._waiting_since = 0.0
303 def _emit(self, state: EngineDisablement, retry: bool = False):
304 message = (
305 f"{state.engine}: cooldown ended; available for retry"
306 if retry
307 else state.describe()
308 )
309 if retry:
310 logger.info(message)
311 else:
312 logger.warning(message)
313 if self._on_change:
314 try:
315 self._on_change(
316 message,
317 None,
318 {
319 "phase": "engine_availability",
320 "type": "info" if retry else "warning",
321 "engine": state.engine,
322 "status": "retry_available" if retry else "disabled",
323 "reason": state.failure.reason,
324 "disabled_at": state.disabled_at,
325 "retry_at": state.retry_at,
326 },
327 )
328 except Exception:
329 logger.exception("Could not report search engine availability")
331 def disable(
332 self, engine: str, failure: SearchFailure
333 ) -> EngineDisablement | UnrecordedFailure:
334 if not engine:
335 # A missing engine name maps to every non-search tool in
336 # engine_for_tool; recording it would hide fetch/delegation.
337 logger.warning("Ignoring availability failure without an engine")
338 return UnrecordedFailure(failure)
339 with self._lock:
340 current = self._disabled.get(engine)
341 if (
342 current is not None
343 and (
344 current.deadline is None or self._clock() < current.deadline
345 )
346 # A run-long failure (e.g. HTTP 401 from a search already
347 # in flight) upgrades a cooldown instead of being dropped.
348 and not (
349 failure.cooldown_seconds is None
350 and current.deadline is not None
351 )
352 ):
353 return current
354 now = self._wall_clock()
355 delay = failure.cooldown_seconds
356 if delay is not None:
357 delay = retry_after_seconds(delay, now=now) or 60
358 count = self._consecutive_cooldowns.get(engine, 0) + 1
359 longest = max(self._wait_budget, MAX_TEMPORARY_COOLDOWN_SECONDS)
360 if delay > longest:
361 failure = SearchFailure(
362 f"{failure.reason}; the provider asked to wait "
363 f"longer than {math.ceil(longest)}s",
364 None,
365 )
366 delay = None
367 elif (
368 self._max_consecutive_cooldowns is not None
369 and count > self._max_consecutive_cooldowns
370 ):
371 failure = SearchFailure(
372 f"{failure.reason}; still failing after "
373 f"{count - 1} cooldowns in a row",
374 None,
375 )
376 delay = None
377 else:
378 self._consecutive_cooldowns[engine] = count
379 state = EngineDisablement(
380 engine,
381 failure,
382 now,
383 None if delay is None else now + delay,
384 None if delay is None else self._clock() + delay,
385 )
386 self._disabled[engine] = state
387 self._emit(state)
388 return state
390 def disabled(self) -> dict[str, EngineDisablement]:
391 expired = []
392 with self._lock:
393 now = self._clock()
394 for engine, state in list(self._disabled.items()):
395 if state.deadline is not None and now >= state.deadline:
396 expired.append(self._disabled.pop(engine))
397 states = dict(self._disabled)
398 for state in expired:
399 self._emit(state, retry=True)
400 return states
402 def note_cancelled(self) -> None:
403 """Record that this run was cancelled, for every thread sharing it.
405 The run's termination flag is single-use: the first check to see it
406 clears it while cleaning up, and so does the cleanup after the lead
407 agent stops. In the LangGraph strategy that check is the research
408 progress callback, which records every cancel it raises, on any
409 thread; drain loops, subagent workers and cooldown waits also record
410 the cancels they see. Parallel drain loops, cooldown waits and
411 running subagents read this record, so they stop as well.
412 """
413 self._cancelled.set()
415 def cancelled(self) -> bool:
416 """Whether a cancel has been recorded (``note_cancelled``)."""
417 return self._cancelled.is_set()
419 def now(self) -> float:
420 """Current time on the clock that cooldown deadlines use."""
421 return self._clock()
423 def record_success(self, engine: str) -> None:
424 """The engine answered a search; its cooldown streak starts over."""
425 with self._lock:
426 self._consecutive_cooldowns.pop(engine, None)
428 @contextmanager
429 def waiting(self):
430 """Charge the time spent inside to the run's cooldown-wait budget."""
431 with self._lock:
432 if self._waiters == 0:
433 self._waiting_since = self._clock()
434 self._waiters += 1
435 try:
436 yield
437 finally:
438 with self._lock:
439 self._waiters -= 1
440 if self._waiters == 0:
441 self._waited += self._clock() - self._waiting_since
443 def wait_budget_remaining(self) -> float:
444 """Seconds of cooldown waiting this run may still spend."""
445 with self._lock:
446 waited = self._waited
447 if self._waiters:
448 waited += self._clock() - self._waiting_since
449 return max(0.0, self._wait_budget - waited)
451 def end_cooldowns(self, engines) -> None:
452 """Disable ``engines`` still cooling down for the rest of the run.
454 Used once the run's cooldown-wait budget is spent: waiting again
455 is no longer an option, so the usual "no research tools" handling
456 applies instead of an endless wait-and-retry cycle.
457 """
458 for engine in engines:
459 state = self.state(engine)
460 if state is not None and state.deadline is not None:
461 self.disable(
462 engine,
463 SearchFailure(
464 f"{state.failure.reason}; cooldown exceeds this "
465 "research's remaining search wait budget",
466 None,
467 ),
468 )
470 def seconds_until_retry(self, engines) -> float | None:
471 """Seconds until the soonest of ``engines`` leaves its cooldown.
473 Only finite cooldowns count; None when none of them has one.
474 """
475 with self._lock:
476 now = self._clock()
477 remaining = [
478 state.deadline - now
479 for engine, state in self._disabled.items()
480 if engine in engines and state.deadline is not None
481 ]
482 return max(0.0, min(remaining)) if remaining else None
484 def note_metrics_name(self, engine: str, name: str) -> None:
485 """Remember the name search metrics record for ``engine``.
487 Metrics name engines after their class (``semanticscholar``), not
488 their configuration key (``semantic_scholar``). Skipped searches are
489 recorded under the same name so they group with real ones.
490 """
491 with self._lock:
492 self._metrics_names[engine] = name
494 def metrics_name(self, engine: str) -> str:
495 with self._lock:
496 return self._metrics_names.get(engine, engine)
498 def state(self, engine: str) -> EngineDisablement | None:
499 """The engine's current disablement, or None when it may be searched."""
500 return self.disabled().get(engine)
503def engine_failure(engine) -> SearchFailure | None:
504 """Read explicit engine health, without mistaking empty results for failure."""
505 failure = getattr(engine, "last_search_failure", None)
506 if isinstance(failure, SearchFailure):
507 return failure
508 if getattr(engine, "_is_available", None) is False: 508 ↛ 509line 508 didn't jump to line 509 because the condition on line 508 was never true
509 return SearchFailure("Search engine failed its availability check")
510 return None
513def engine_for_tool(name: str, primary_engine: str) -> str | None:
514 if name == "fetch_archive":
515 return "wayback"
516 if name == "web_search":
517 return primary_engine
518 if name.startswith("search_"):
519 return name.removeprefix("search_")
520 return None
523def search_engines_for_tools(names, primary_engine) -> set[str]:
524 """Providers that discover URLs; archive fetching only reads known URLs."""
525 engines = {
526 engine_for_tool(name, primary_engine)
527 for name in names
528 if name == "web_search" or name.startswith("search_")
529 }
530 return {engine for engine in engines if engine is not None}
533def search_exhausted(names, primary_engine, disabled) -> bool:
534 """The toolbox had query search tools, and every one is disabled."""
535 engines = search_engines_for_tools(names, primary_engine)
536 return bool(engines) and engines <= disabled.keys()
539def available_tool_names(names, primary_engine, disabled) -> list[str]:
540 """Filter an already policy-approved toolbox, including empty delegation.
542 Delegation is dropped once availability has removed every search tool:
543 its subagents would be left with fetch tools alone and no URLs
544 to read. (A toolbox configured without search is left as it is.)
545 """
546 available = [
547 name
548 for name in names
549 if (engine := engine_for_tool(name, primary_engine)) is None
550 or engine not in disabled
551 ]
552 if search_exhausted(names, primary_engine, disabled):
553 available = [name for name in available if name != "research_subtopic"]
554 if available == ["research_subtopic"]: 554 ↛ 555line 554 didn't jump to line 555 because the condition on line 554 was never true
555 return []
556 return available