Coverage for src/local_deep_research/utilities/research_content_cache.py: 82%
234 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""Single shared content cache for an entire research run.
3Every consumer that needs the full text of a webpage, a RAG chunk, a
4library document — whether it is the ``web_search`` tool (via
5``FullSearchResults`` / ``BaseSearchEngine._get_full_content``), the
6``fetch_content`` tool (via ``ContentFetcher`` + the summary-LLM prompt
7construction), or the library / collection RAG engines — resolves through
8this one cache. The initial quick-report phase, the detailed-report
9subsections and all parallel subagents therefore share previously fetched
10content instead of re-fetching the same URLs on the spot.
12Lifetime: one instance per research run, owned by ``AdvancedSearchSystem``
13(``system.content_cache``) and threaded explicitly into strategies, tools
14and engines. Deep funnels that cannot receive an explicit reference
15(``batch_fetch_and_extract``, ``ContentFetcher.fetch``,
16``resolve_library_document``) fall back to the research-scoped registry
17below, keyed by the ``research_id`` carried in the thread-local search
18context — so freshly created per-call engines still hit the same cache.
20When no cache is bound (no explicit instance and no research context, e.g.
21legacy unit tests) all helpers return ``None`` and callers fall back to
22direct fetching, preserving historical behaviour.
23"""
25from __future__ import annotations
27import threading
28from typing import Any, Callable
31def normalize_collection_scope(collection_id: Any) -> str | None:
32 """Normalise a collection identifier to its ``collection_<id>`` form.
34 Returns ``None`` for missing/empty values. A value already in
35 ``collection_<id>`` form is returned unchanged so ``"coll-B"`` and
36 ``"collection_coll-B"`` collapse to one cache namespace (matching the
37 ``DocumentChunk.collection_name`` normalisation in
38 ``LibraryRAGSearchEngine._get_full_content``). Human-readable names
39 that do not look like ``collection_*`` still get a distinct namespace
40 (under-merge, never cross-contamination).
41 """
42 if collection_id is None or isinstance(collection_id, bool):
43 return None
44 if not isinstance(collection_id, str): 44 ↛ 45line 44 didn't jump to line 45 because the condition on line 44 was never true
45 collection_id = str(collection_id)
46 stripped = collection_id.strip()
47 if not stripped: 47 ↛ 48line 47 didn't jump to line 48 because the condition on line 47 was never true
48 return None
49 if stripped.startswith("collection_"):
50 return stripped
51 return f"collection_{stripped}"
54def _normalize_chunk_display(display: str) -> str:
55 """Merge the two chunk spellings to one cache identity.
57 ``/library/document/<id>/chunks#chunk-N`` (emitted by
58 ``build_chunk_anchor_url``) and ``/library/document/<id>#chunk-N``
59 (accepted by the resolver/tests) denote the same chunk text. Strip a
60 trailing ``/chunks`` view when a chunk fragment is present so both
61 share one entry (under-merge only, no corruption).
62 """
63 if "#chunk-" not in display: 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true
64 return display
65 path, _, fragment = display.partition("#")
66 if path.endswith("/chunks"):
67 path = path[: -len("/chunks")]
68 return f"{path}#{fragment}"
69 return display
72def content_cache_key(url: Any, collection_id: Any = None) -> str:
73 """Normalise *url* into a stable cache key.
75 Web URLs collapse tracking params / default ports via
76 ``content_web_cache_key`` (NOT the bibliography ``canonical_url_key``:
77 the path keeps its trailing slash, userinfo is preserved, and
78 fragments are preserved for client-routed SPAs) while
79 ``https://example.com/p?utm_x=1`` and ``https://example.com/p`` share
80 one entry (accepting potential over-merge on sites that vary content per
81 tracking param). Library routes keep their chunk anchor (``#chunk-<n>``)
82 via ``library_display_url`` because a chunk and its parent document
83 are *different* content and must not collide. Anything unparseable
84 falls back to the stripped raw string.
86 Library *chunk* URLs are additionally scoped by *collection_id*:
87 the same ``(document, chunk_index)`` can hold different text in two
88 collections (different chunk settings), so ``collection_A#chunk-3``
89 and ``collection_B#chunk-3`` must not share an entry. Document-level
90 library URLs (``Document.text_content`` is collection-independent)
91 and web URLs ignore *collection_id*.
92 """
93 if not isinstance(url, str):
94 return f"raw:{url!r}"
95 stripped = url.strip()
96 if not stripped:
97 return "raw:"
98 try:
99 from .url_utils import content_web_cache_key, library_display_url
100 except Exception: # pragma: no cover - defensive
101 return f"raw:{stripped}"
102 try:
103 display = library_display_url(stripped)
104 except Exception: # pragma: no cover - defensive
105 display = None
106 if display is not None:
107 if "#chunk-" in display:
108 normalised = _normalize_chunk_display(display)
109 scope = normalize_collection_scope(collection_id)
110 if scope is not None:
111 return f"lib:{normalised}|collection:{scope}"
112 return f"lib:{normalised}|collection:__unscoped__"
113 return f"lib:{display}"
114 try:
115 canon = content_web_cache_key(stripped)
116 except Exception: # pragma: no cover - defensive
117 canon = ""
118 if canon: 118 ↛ 120line 118 didn't jump to line 120 because the condition on line 118 was always true
119 return f"web:{canon}"
120 return f"raw:{stripped}"
123class ResearchContentCache:
124 """Thread-safe in-memory store of fetched page / chunk text.
126 Entries hold the *fullest* fetched content (before per-caller
127 truncation). Callers slice to their own ``max_length`` on the way out,
128 so a ``fetch_content`` call capped at 10k chars and a ``web_search``
129 full-content call share one entry without data loss: the longer fetch
130 wins on ``put``.
132 Memory is bounded: at most ``max_entries`` entries and
133 ``max_bytes`` of content bytes are retained (evict-oldest/FIFO).
134 A single oversized page is still stored (it evicts everything else
135 rather than being dropped) so callers never lose the entry they just
136 fetched.
137 """
139 # Defaults keep a full research run (quick phase + subsections +
140 # subagents) in memory while capping multi-user server growth:
141 # 200 x ~100KB typical pages ~= 20MB; worst case 32MB per run.
142 DEFAULT_MAX_ENTRIES = 200
143 DEFAULT_MAX_BYTES = 32 * 1024 * 1024
144 # Cap for the SSRF-verdict memo (per-run, FIFO). Each entry is a
145 # small bool keyed by (url, allow_private, block_link_local); 1000
146 # covers a full run's distinct URLs without unbounded growth.
147 MAX_VALIDATION_MEMO_ENTRIES = 1000
149 def __init__(
150 self,
151 max_entries: int = DEFAULT_MAX_ENTRIES,
152 max_bytes: int = DEFAULT_MAX_BYTES,
153 ) -> None:
154 self._lock = threading.RLock()
155 # key -> {"content": str, "title": str|None, "url": str}
156 self._entries: dict[str, dict[str, Any]] = {}
157 # key -> threading.Event for in-flight fetches (thundering-herd guard)
158 self._inflight: dict[str, threading.Event] = {}
159 self._max_entries = max(1, int(max_entries))
160 self._max_bytes = max(1, int(max_bytes))
161 self._total_bytes = 0
162 self._hits = 0
163 self._misses = 0
164 self._puts = 0
165 self._evictions = 0
166 # (url, allow_private, block_link_local) -> True. Per-run memo
167 # for the SSRF grant gates so repeat cache hits do not pay a
168 # getaddrinfo per read (and stay servable when DNS drops
169 # mid-run after a successful verdict). Only True verdicts are
170 # pinned: validate_url converts a transient DNS failure
171 # (socket.gaierror) into a False verdict rather than raising,
172 # so pinning False would block every later fetch of that URL
173 # for the rest of the run after one blip — a False verdict only
174 # skips the cache fast path and the fetcher/downloader
175 # re-validates on its own, so recomputing False never blocks
176 # recovery.
177 self._validation_memo: dict[tuple[str, bool, bool], bool] = {}
179 # -- basic ops ------------------------------------------------------
180 def get(self, url: str, collection_id: Any = None) -> dict[str, Any] | None:
181 key = content_cache_key(url, collection_id)
182 with self._lock:
183 entry = self._entries.get(key)
184 if entry is not None:
185 self._hits += 1
186 return dict(entry)
187 self._misses += 1
188 return None
190 def _evict_if_needed(self) -> None:
191 """Evict oldest entries until within entry/byte caps (caller holds lock)."""
192 while (
193 len(self._entries) > self._max_entries
194 or self._total_bytes > self._max_bytes
195 ) and len(self._entries) > 1:
196 oldest_key = next(iter(self._entries))
197 oldest = self._entries.pop(oldest_key)
198 self._total_bytes -= len(oldest.get("content", "") or "")
199 if self._total_bytes < 0: 199 ↛ 200line 199 didn't jump to line 200 because the condition on line 199 was never true
200 self._total_bytes = 0
201 self._evictions += 1
203 def put(
204 self,
205 url: str,
206 content: str,
207 title: str | None = None,
208 url_out: str | None = None,
209 author: str | None = None,
210 published_date: Any | None = None,
211 collection_id: Any = None,
212 ) -> None:
213 """Store *content* for *url*, keeping the longest version on conflict.
215 Only full pre-truncation text must be stored here — callers slice
216 to their own ``max_length`` (plus any truncation marker) on the way
217 out. Storing truncated display output would let a
218 ``10_000 + marker`` slice overwrite the complete ``10_001``-char
219 page under longest-wins and lose the tail for later callers with
220 larger limits.
222 Library chunk URLs are namespaced by *collection_id* (see
223 :func:`content_cache_key`): the same chunk index can hold
224 different text per collection.
225 """
226 if not isinstance(content, str): 226 ↛ 227line 226 didn't jump to line 227 because the condition on line 226 was never true
227 return
228 key = content_cache_key(url, collection_id)
229 with self._lock:
230 existing = self._entries.get(key)
231 if existing is not None and len(existing.get("content", "")) >= len(
232 content
233 ):
234 return
235 if existing is not None:
236 self._total_bytes -= len(existing.get("content", "") or "")
237 # Re-insert as newest so FIFO eviction drops truly old URLs.
238 self._entries.pop(key, None)
239 entry: dict[str, Any] = {
240 "content": content,
241 "title": title,
242 "url": url_out if url_out is not None else url,
243 }
244 if author is not None: 244 ↛ 245line 244 didn't jump to line 245 because the condition on line 244 was never true
245 entry["author"] = author
246 if published_date is not None: 246 ↛ 247line 246 didn't jump to line 247 because the condition on line 246 was never true
247 entry["published_date"] = published_date
248 self._entries[key] = entry
249 self._total_bytes += len(content)
250 self._puts += 1
251 self._evict_if_needed()
253 def get_or_fetch(
254 self,
255 url: str,
256 fetch_fn: Callable[[], dict[str, Any] | None],
257 *,
258 wait_timeout: float = 120.0,
259 collection_id: Any = None,
260 ) -> dict[str, Any] | None:
261 """Return the cached entry for *url* or fetch it exactly once.
263 Concurrent callers for the same URL coalesce: the first becomes the
264 fetcher, the rest wait on an event and then read the cache. If the
265 fetcher fails (exception or ``None``) waiters fall through and try
266 their own fetch instead of caching the failure. A waiter whose owner
267 does not finish within ``wait_timeout`` raises ``TimeoutError`` without
268 cancelling the owner; later callers can still reuse its result.
269 Results marked ``content_is_complete=False`` are returned without
270 caching, so a metadata fallback cannot prevent a later full fetch.
272 The result is :meth:`put` *before* waiters are woken — waking first
273 would let a waiter re-enter as a second fetcher and double-fetch.
275 Library chunk URLs are namespaced by *collection_id* (see
276 :func:`content_cache_key`).
277 """
278 key = content_cache_key(url, collection_id)
279 while True:
280 with self._lock:
281 entry = self._entries.get(key)
282 if entry is not None:
283 self._hits += 1
284 return dict(entry)
285 self._misses += 1
286 event = self._inflight.get(key)
287 if event is None:
288 event = threading.Event()
289 self._inflight[key] = event
290 break
291 if not event.wait(timeout=wait_timeout):
292 # Keep the original owner's claim: a timed-out waiter must
293 # neither recurse forever nor release another thread's fetch.
294 raise TimeoutError("Timed out waiting for shared content")
295 # The owner completed without caching a result. Retry election
296 # iteratively so repeated failures never grow the call stack.
297 try:
298 result = fetch_fn()
299 if result is None:
300 return None
301 content = (
302 result.get("content") if isinstance(result, dict) else None
303 )
304 if (
305 isinstance(content, str)
306 and result.get("content_is_complete") is not False
307 ):
308 # Cache BEFORE waking waiters so they observe the fresh
309 # entry. The finally below wakes waiters; put failures
310 # still wake via finally so waiters retry instead of
311 # hanging on a dead event.
312 self.put(
313 url,
314 content,
315 title=result.get("title")
316 if isinstance(result, dict)
317 else None,
318 url_out=result.get("url")
319 if isinstance(result, dict)
320 else None,
321 author=result.get("author")
322 if isinstance(result, dict)
323 else None,
324 published_date=result.get("published_date")
325 if isinstance(result, dict)
326 else None,
327 collection_id=collection_id,
328 )
329 with self._lock:
330 cached = self._entries.get(key)
331 if cached is not None:
332 return dict(cached)
333 return result
334 finally:
335 # Always release the in-flight slot and wake waiters — even
336 # for BaseException (cancellation/KeyboardInterrupt) and put
337 # failures — or a later caller would block until wait_timeout
338 # on a dead event. Store-before-wake holds because put runs
339 # inside try before this finally signals the event.
340 with self._lock:
341 self._inflight.pop(key, None)
342 event.set()
344 # -- SSRF-verdict memo (per-run, defense-in-depth) ------------------
345 @staticmethod
346 def _validation_memo_key(
347 url: Any, allow_private_ips: bool, block_link_local: bool
348 ) -> tuple[str, bool, bool] | None:
349 if not isinstance(url, str): 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true
350 return None
351 stripped = url.strip()
352 if not stripped: 352 ↛ 353line 352 didn't jump to line 353 because the condition on line 352 was never true
353 return None
354 return (stripped, bool(allow_private_ips), bool(block_link_local))
356 def get_validation(
357 self, url: str, allow_private_ips: bool, block_link_local: bool
358 ) -> bool | None:
359 """Return the memoized SSRF verdict, or None when not memoized."""
360 key = self._validation_memo_key(
361 url, allow_private_ips, block_link_local
362 )
363 if key is None: 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true
364 return None
365 with self._lock:
366 return self._validation_memo.get(key)
368 def set_validation(
369 self,
370 url: str,
371 allow_private_ips: bool,
372 block_link_local: bool,
373 allowed: bool,
374 ) -> None:
375 """Memoize a successful SSRF verdict (FIFO-bounded, per-run).
377 Only True is stored — see :meth:`validated` for why False is
378 never pinned.
379 """
380 if not allowed:
381 return
382 key = self._validation_memo_key(
383 url, allow_private_ips, block_link_local
384 )
385 if key is None:
386 return
387 with self._lock:
388 if key not in self._validation_memo:
389 while (
390 len(self._validation_memo)
391 >= self.MAX_VALIDATION_MEMO_ENTRIES
392 ):
393 try:
394 self._validation_memo.pop(
395 next(iter(self._validation_memo))
396 )
397 except StopIteration:
398 break
399 self._validation_memo[key] = True
401 def validated(
402 self,
403 url: str,
404 allow_private_ips: bool,
405 block_link_local: bool,
406 compute: Callable[[], bool],
407 ) -> bool:
408 """Return memoized success or compute, pinning only True.
410 Thread-safe: concurrent callers for the same key may compute in
411 parallel (DNS outside the lock); the first True wins and stays
412 for the run like ``policy._dns_cache``. False is never stored:
413 ``validate_url`` maps a transient DNS failure (socket.gaierror)
414 to False instead of raising, so pinning it would refuse every
415 later fetch of that URL until the run ends even after DNS heals.
416 Recomputing False costs one resolution per refused hit (bounded,
417 perf-only) and keeps the fetcher/downloader recovery path open.
418 """
419 key = self._validation_memo_key(
420 url, allow_private_ips, block_link_local
421 )
422 if key is None: 422 ↛ 423line 422 didn't jump to line 423 because the condition on line 422 was never true
423 try:
424 return bool(compute())
425 except Exception:
426 return False
427 with self._lock:
428 if key in self._validation_memo:
429 return True
430 try:
431 verdict = bool(compute())
432 except Exception:
433 return False
434 if not verdict:
435 return False
436 with self._lock:
437 if key not in self._validation_memo: 437 ↛ 449line 437 didn't jump to line 449 because the condition on line 437 was always true
438 while ( 438 ↛ 442line 438 didn't jump to line 442 because the condition on line 438 was never true
439 len(self._validation_memo)
440 >= self.MAX_VALIDATION_MEMO_ENTRIES
441 ):
442 try:
443 self._validation_memo.pop(
444 next(iter(self._validation_memo))
445 )
446 except StopIteration:
447 break
448 self._validation_memo[key] = True
449 return True
451 def clear(self) -> None:
452 with self._lock:
453 self._entries.clear()
454 self._inflight.clear()
455 self._validation_memo.clear()
456 self._total_bytes = 0
457 self._hits = 0
458 self._misses = 0
459 self._puts = 0
460 self._evictions = 0
462 def stats(self) -> dict[str, int]:
463 with self._lock:
464 return {
465 "hits": self._hits,
466 "misses": self._misses,
467 "puts": self._puts,
468 "entries": len(self._entries),
469 "evictions": self._evictions,
470 "total_bytes": self._total_bytes,
471 "validation_memo": len(self._validation_memo),
472 }
474 def __len__(self) -> int:
475 with self._lock:
476 return len(self._entries)
479# -- research-scoped registry -------------------------------------------
480# Maps research_id -> ResearchContentCache. Owned by AdvancedSearchSystem
481# (one cache per research run); deep funnels resolve it via the
482# thread-local search context so per-call engines share it without
483# signature changes across dozens of files.
485_registry_lock = threading.RLock()
486_registry: dict[str, ResearchContentCache] = {}
489def bind_cache_to_research(
490 research_id: str | None, cache: ResearchContentCache | None
491) -> None:
492 """Bind *cache* to *research_id* (no-op for empty ids / None cache)."""
493 if not research_id or cache is None:
494 return
495 with _registry_lock:
496 _registry[research_id] = cache
499def unbind_cache_from_research(
500 research_id: str | None, *, cache: ResearchContentCache | None = None
501) -> None:
502 """Unbind a run, optionally only when *cache* still owns its binding."""
503 if not research_id: 503 ↛ 504line 503 didn't jump to line 504 because the condition on line 503 was never true
504 return
505 with _registry_lock:
506 if cache is None or _registry.get(research_id) is cache:
507 _registry.pop(research_id, None)
510def get_cache_for_research(
511 research_id: str | None,
512) -> ResearchContentCache | None:
513 if not research_id:
514 return None
515 with _registry_lock:
516 return _registry.get(research_id)
519def get_current_cache() -> ResearchContentCache | None:
520 """Resolve the active research's cache via the thread-local context."""
521 try:
522 from .thread_context import get_search_context
523 except Exception: # pragma: no cover - defensive
524 return None
525 try:
526 ctx = get_search_context()
527 except Exception: # pragma: no cover - defensive
528 return None
529 if not isinstance(ctx, dict):
530 return None
531 research_id = ctx.get("research_id")
532 if not research_id or not isinstance(research_id, str):
533 # Non-string ids (int legacy) are still valid scopes.
534 if research_id is None or isinstance(research_id, bool): 534 ↛ 536line 534 didn't jump to line 536 because the condition on line 534 was always true
535 return None
536 research_id = str(research_id)
537 return get_cache_for_research(research_id)
540def resolve_cache(
541 explicit: ResearchContentCache | None,
542) -> ResearchContentCache | None:
543 """Prefer *explicit*, else the registry-bound cache for this thread."""
544 if explicit is not None:
545 return explicit
546 return get_current_cache()
549__all__ = [
550 "ResearchContentCache",
551 "bind_cache_to_research",
552 "content_cache_key",
553 "get_cache_for_research",
554 "get_current_cache",
555 "normalize_collection_scope",
556 "resolve_cache",
557 "unbind_cache_from_research",
558]