Coverage for astrocyte/recall/proxy.py: 91%

390 statements  

« prev     ^ index     » next       coverage.py v7.15.0, created at 2026-10-29 22:10 +0000

1"""HTTP proxy recall — fetch remote hits and merge with local RRF (M4.1).""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import ipaddress 

7import json 

8import logging 

9import os 

10import re 

11import socket 

12import time 

13from datetime import datetime, timezone 

14from typing import TYPE_CHECKING, Any 

15from urllib.parse import ParseResult, parse_qsl, quote, urlparse, urlunparse 

16 

17import httpx 

18 

19from astrocyte.config import SourceConfig 

20from astrocyte.policy.observability import MetricsCollector, span, timed 

21from astrocyte.types import MemoryHit, Metadata 

22 

23if TYPE_CHECKING: 

24 from astrocyte.config import AstrocyteConfig 

25 

26logger = logging.getLogger(__name__) 

27 

28_DEFAULT_TIMEOUT = 15.0 

29 

30# JSON body placeholders (POST) — resolved to ``query`` / ``bank_id`` strings 

31PLACE_QUERY = "__astrocyte.query__" 

32PLACE_BANK = "__astrocyte.bank_id__" 

33 

34_COUNTER = "astrocyte_proxy_recall_total" 

35_HIST = "astrocyte_proxy_recall_duration_seconds" 

36 

37 

38def _expand_proxy_url(template: str, query: str) -> str: 

39 if "{query}" in template: 

40 return template.replace("{query}", quote(query, safe="")) 

41 sep = "&" if "?" in template else "?" 

42 return f"{template}{sep}q={quote(query, safe='')}" 

43 

44 

45def _forbidden_proxy_target_ip_obj(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: 

46 """True if this address must not be used for outbound proxy recall (SSRF mitigation).""" 

47 return bool( 

48 ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved or ip.is_multicast or ip.is_unspecified 

49 ) 

50 

51 

52def _unsafe_literal_ip(host: str) -> bool: 

53 """True if *host* parses as an IPv4/IPv6 address that must not be used for proxy recall.""" 

54 try: 

55 ip = ipaddress.ip_address(host) 

56 except ValueError: 

57 return False 

58 return _forbidden_proxy_target_ip_obj(ip) 

59 

60 

61def validate_proxy_recall_url(url: str) -> None: 

62 """Reject URLs that enable SSRF (private/loopback/metadata-ranged IPs, non-HTTP(S), no host). 

63 

64 Call this on the **fully expanded** request URL (including encoded user ``query`` fragments). 

65 Hostnames still require :func:`validate_proxy_recall_dns` before HTTP (rebinding-safe check). 

66 """ 

67 if not (url or "").strip(): 

68 raise ValueError("proxy recall URL is empty") 

69 parsed = urlparse(url.strip()) 

70 if parsed.scheme not in ("http", "https"): 

71 raise ValueError(f"proxy recall URL scheme not allowed: {parsed.scheme!r}") 

72 host = parsed.hostname 

73 if not host: 

74 raise ValueError("proxy recall URL has no host") 

75 h = host.lower().rstrip(".") 

76 if h == "localhost" or h.endswith(".localhost"): 

77 raise ValueError("proxy recall URL must not target localhost") 

78 if _unsafe_literal_ip(h): 

79 raise ValueError( 

80 "proxy recall URL must not target loopback, private, link-local, or reserved addresses", 

81 ) 

82 

83 

84def _sync_dns_validate_and_first_public_ip( 

85 host: str, 

86 port: int, 

87) -> ipaddress.IPv4Address | ipaddress.IPv6Address: 

88 """Resolve *host*, reject if any address is forbidden, return the first (OS order) for pinning.""" 

89 try: 

90 infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM) 

91 except socket.gaierror as e: 

92 raise ValueError(f"proxy recall DNS resolution failed for {host!r}: {e}") from e 

93 if not infos: 

94 raise ValueError(f"proxy recall DNS returned no addresses for {host!r}") 

95 picked: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = [] 

96 for _fam, _socktype, _proto, _canon, sockaddr in infos: 

97 addr_s = sockaddr[0] 

98 try: 

99 ip = ipaddress.ip_address(addr_s) 

100 except ValueError: 

101 continue 

102 if _forbidden_proxy_target_ip_obj(ip): 

103 raise ValueError( 

104 f"proxy recall DNS resolved to a forbidden address ({addr_s!r} for host {host!r})", 

105 ) 

106 picked.append(ip) 

107 if not picked: 

108 raise ValueError(f"proxy recall DNS returned no usable addresses for {host!r}") 

109 return picked[0] 

110 

111 

112def _is_literal_ip_host(host: str) -> bool: 

113 """True if *host* is an IPv4/IPv6 literal (no DNS name).""" 

114 try: 

115 ipaddress.ip_address(host) 

116 except ValueError: 

117 return False 

118 return True 

119 

120 

121def _proxy_recall_host_header_value(original_hostname: str, port: int, scheme: str) -> str: 

122 """``Host`` header for the original server name (include port when not default).""" 

123 sch = (scheme or "").lower() 

124 default = 443 if sch == "https" else 80 

125 h = original_hostname.lower().rstrip(".") 

126 if port != default: 

127 return f"{h}:{port}" 

128 return h 

129 

130 

131def _rebuild_request_url_pinned_to_ip( 

132 parsed: ParseResult, 

133 *, 

134 pinned: ipaddress.IPv4Address | ipaddress.IPv6Address, 

135 port: int, 

136) -> str: 

137 """Replace authority with pinned IP so the TCP/TLS layer cannot re-resolve differently.""" 

138 if isinstance(pinned, ipaddress.IPv6Address): 

139 hostpart = f"[{pinned.compressed}]" 

140 else: 

141 hostpart = str(pinned) 

142 sch = (parsed.scheme or "").lower() 

143 default = 443 if sch == "https" else 80 

144 if port != default: 

145 netloc = f"{hostpart}:{port}" 

146 else: 

147 netloc = hostpart 

148 return urlunparse((parsed.scheme, netloc, parsed.path, parsed.params, parsed.query, parsed.fragment)) 

149 

150 

151def _httpx_url_and_query_params(url: str) -> tuple[str, list[tuple[str, str]] | None]: 

152 """Split *url* into a path-only URL for the httpx *url* argument and query pairs for *params*. 

153 

154 User-controlled search text must not be concatenated into the URL string passed to httpx 

155 (SSRF static analysis + clearer separation of authority vs query). 

156 """ 

157 p = urlparse((url or "").strip()) 

158 without_query = urlunparse((p.scheme, p.netloc, p.path, p.params, "", p.fragment)) 

159 if not p.query: 

160 return without_query, None 

161 pairs = parse_qsl(p.query, keep_blank_values=True) 

162 if not pairs: 

163 return without_query, None 

164 return without_query, list(pairs) 

165 

166 

167async def validate_proxy_recall_dns(host: str, port: int) -> None: 

168 """Async-safe DNS check: resolve *host* in a worker thread, then validate all addresses.""" 

169 await asyncio.to_thread(_sync_dns_validate_and_first_public_ip, host, port) 

170 

171 

172def auth_with_oauth_cache_namespace( 

173 auth: dict[str, str | int | float | bool | None] | None, 

174 source_id: str, 

175) -> dict[str, str | int | float | bool | None] | None: 

176 """Attach ``_oauth_cache_id`` so OAuth token caches do not collide across proxy sources.""" 

177 if not auth: 

178 return None 

179 return {**auth, "_oauth_cache_id": source_id} 

180 

181 

182async def build_proxy_headers( 

183 auth: dict[str, str | int | float | bool | None] | None, 

184) -> dict[str, str]: 

185 """Build HTTP headers (Bearer, API key, OAuth2 client_credentials / refresh, optional static ``headers``).""" 

186 out: dict[str, str] = {} 

187 if not auth: 

188 return out 

189 extra = auth.get("headers") 

190 if isinstance(extra, dict): 

191 for k, v in extra.items(): 

192 if isinstance(v, (str, int, float, bool)): 

193 out[str(k)] = str(v) 

194 t = (str(auth.get("type") or "")).strip().lower() 

195 grant = (str(auth.get("grant_type") or "")).strip().lower() 

196 if t in ("oauth2", "oauth2_client_credentials") and grant in ("", "client_credentials"): 

197 from astrocyte.recall.oauth import fetch_oauth2_client_credentials_token 

198 

199 token = await fetch_oauth2_client_credentials_token(auth) 

200 out["Authorization"] = f"Bearer {token}" 

201 elif t == "oauth2_refresh" or (t in ("oauth2",) and grant == "refresh_token"): 

202 from astrocyte.recall.oauth import fetch_oauth2_refresh_access_token 

203 

204 token = await fetch_oauth2_refresh_access_token(auth) 

205 out["Authorization"] = f"Bearer {token}" 

206 elif t == "bearer": 

207 token = auth.get("token") 

208 if token is not None and str(token).strip(): 

209 out["Authorization"] = f"Bearer {token}" 

210 elif t == "api_key": 

211 header_name = str(auth.get("header") or "X-API-Key") 

212 if not re.match(r"^[A-Za-z0-9\-]+$", header_name): 

213 raise ValueError(f"Invalid header name in proxy auth config: {header_name!r}") 

214 val = auth.get("value") if auth.get("value") is not None else auth.get("token") 

215 if val is not None and str(val).strip(): 

216 val_str = str(val) 

217 if "\r" in val_str or "\n" in val_str: 

218 raise ValueError("Header value contains CRLF characters (possible injection)") 

219 out[header_name] = val_str 

220 return out 

221 

222 

223def _deep_replace_placeholders(obj: Any, query: str, bank_id: str) -> Any: 

224 if obj == PLACE_QUERY: 

225 return query 

226 if obj == PLACE_BANK: 

227 return bank_id 

228 if isinstance(obj, dict): 

229 return {k: _deep_replace_placeholders(v, query, bank_id) for k, v in obj.items()} 

230 if isinstance(obj, list): 

231 return [_deep_replace_placeholders(x, query, bank_id) for x in obj] 

232 return obj 

233 

234 

235def _resolve_post_json(source: SourceConfig, query: str, bank_id: str) -> dict[str, Any]: 

236 rb = source.recall_body 

237 if rb is None: 

238 return {"query": query, "bank_id": bank_id} 

239 if isinstance(rb, dict): 

240 resolved = _deep_replace_placeholders(rb, query, bank_id) 

241 if isinstance(resolved, dict): 

242 return resolved 

243 return {"query": query, "bank_id": bank_id} 

244 if isinstance(rb, str): 

245 try: 

246 data = json.loads(rb) 

247 except json.JSONDecodeError: 

248 return {"query": query, "bank_id": bank_id} 

249 resolved = _deep_replace_placeholders(data, query, bank_id) 

250 if isinstance(resolved, dict): 

251 return resolved 

252 return {"query": query, "bank_id": bank_id} 

253 return {"query": query, "bank_id": bank_id} 

254 

255 

256#: Provenance fields a remote row may carry, mapped onto reserved metadata keys. 

257#: Reserved keys are written last, so a source cannot spoof them through its own 

258#: ``metadata`` object. 

259_PROVENANCE_FIELDS: dict[str, tuple[str, ...]] = { 

260 "_source_url": ("url", "source_url"), 

261 "_source_version": ("version", "etag", "revision"), 

262 "_source_author": ("author",), 

263 "_source_anchor": ("anchor",), 

264} 

265 

266 

267def _parse_when(value: Any) -> datetime | None: 

268 """ISO-8601 string or Unix seconds -> aware UTC datetime; anything else -> None.""" 

269 if isinstance(value, bool): 

270 return None 

271 try: 

272 if isinstance(value, (int, float)): 

273 return datetime.fromtimestamp(float(value), tz=timezone.utc) 

274 if isinstance(value, str) and value.strip(): 

275 dt = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) 

276 return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc) 

277 except (ValueError, OverflowError, OSError): 

278 return None 

279 return None 

280 

281 

282def _row_to_hit(source_id: str, row: dict[str, Any]) -> MemoryHit | None: 

283 """Map one remote row to a ``MemoryHit`` with the dates and provenance a 

284 local hit carries (federated-sources F0b). 

285 

286 Before 2026-10-04 only text, score, flat metadata, tags, id, and fact type 

287 survived, so every federated hit arrived undated and unanchored: the same 

288 class of defect as a rerank that dropped ``occurred_at`` (roadmap §4e). 

289 Fusion ranks by list position, so a missing score never ranked anything; 

290 it gets a neutral placeholder flagged ``_score_missing`` rather than 

291 passing as a measurement. 

292 """ 

293 text = row.get("text") 

294 if not isinstance(text, str) or not text.strip(): 

295 return None 

296 score = row.get("score") 

297 score_missing = isinstance(score, bool) or not isinstance(score, (int, float)) 

298 s = 0.5 if score_missing else float(score) 

299 mid = row.get("memory_id") 

300 meta: Metadata = {} 

301 meta_raw = row.get("metadata") 

302 if isinstance(meta_raw, dict): 

303 for k, v in meta_raw.items(): 

304 if isinstance(v, (str, int, float, bool)) or v is None: 

305 meta[str(k)] = v 

306 for key, names in _PROVENANCE_FIELDS.items(): 

307 meta.pop(key, None) 

308 for name in names: 

309 v = row.get(name) 

310 if isinstance(v, (str, int, float)) and not isinstance(v, bool) and str(v).strip(): 

311 meta[key] = str(v) 

312 break 

313 if score_missing: 

314 meta["_score_missing"] = True 

315 else: 

316 meta.pop("_score_missing", None) 

317 tags = row.get("tags") 

318 tag_list: list[str] | None = None 

319 if isinstance(tags, list): 

320 tag_list = [str(x) for x in tags] 

321 return MemoryHit( 

322 text=text, 

323 score=min(1.0, max(0.0, s)), 

324 fact_type=str(row["fact_type"]) if row.get("fact_type") is not None else None, 

325 metadata=meta or None, 

326 tags=tag_list, 

327 memory_id=str(mid) if mid is not None else None, 

328 source=f"proxy:{source_id}", 

329 occurred_at=_parse_when(row.get("occurred_at")), 

330 # When the source last recorded it: its own retained_at, else the 

331 # document's last modification. 

332 retained_at=_parse_when(row.get("retained_at")) or _parse_when(row.get("updated_at")), 

333 ) 

334 

335 

336def _parse_hits_payload(data: Any) -> list[Any]: 

337 raw_hits = data.get("hits") if isinstance(data, dict) else None 

338 if raw_hits is None and isinstance(data, dict): 

339 raw_hits = data.get("results") 

340 if not isinstance(raw_hits, list): 

341 return [] 

342 return raw_hits 

343 

344 

345def _record_proxy_metrics( 

346 metrics: MetricsCollector | None, 

347 *, 

348 source_id: str, 

349 status: str, 

350 duration_s: float | None, 

351) -> None: 

352 if not metrics: 

353 return 

354 metrics.inc_counter( 

355 _COUNTER, 

356 {"source_id": source_id, "status": status}, 

357 "Proxy recall attempts by source and status", 

358 ) 

359 # Every timed outcome, not only successes: a histogram of successful calls 

360 # alone hides exactly the slow tail (errors after a long wait, deadline 

361 # misses) that p95 is meant to expose. ``status`` stays on the counter, so 

362 # the documented label set (ADR-003) is unchanged. 

363 if duration_s is not None: 

364 metrics.observe_histogram( 

365 _HIST, 

366 duration_s, 

367 {"source_id": source_id}, 

368 "Proxy recall HTTP duration in seconds", 

369 ) 

370 

371 

372async def fetch_proxy_recall_hits( 

373 source_id: str, 

374 source: SourceConfig, 

375 *, 

376 query: str, 

377 bank_id: str, 

378 timeout: float = _DEFAULT_TIMEOUT, 

379 metrics: MetricsCollector | None = None, 

380) -> list[MemoryHit]: 

381 """Call ``source.url`` (GET or POST) and parse JSON ``hits`` / ``results`` arrays.""" 

382 url_t = source.url or "" 

383 if not url_t.strip(): 

384 return [] 

385 

386 method = (source.recall_method or "GET").strip().upper() 

387 if method not in ("GET", "POST"): 

388 logger.warning("proxy source %s: unknown recall_method %r, using GET", source_id, method) 

389 method = "GET" 

390 

391 base_url = url_t.strip() 

392 started = time.monotonic() 

393 

394 with span( 

395 "astrocyte.proxy_recall", 

396 {"source_id": source_id, "method": method, "bank_id": bank_id}, 

397 ): 

398 with timed() as t: 

399 try: 

400 headers = await build_proxy_headers( 

401 auth_with_oauth_cache_namespace(source.auth, source_id), 

402 ) 

403 if method == "POST": 

404 request_url = _expand_proxy_url(base_url, query) if "{query}" in base_url else base_url 

405 body = _resolve_post_json(source, query, bank_id) 

406 else: 

407 request_url = _expand_proxy_url(base_url, query) 

408 body = None 

409 validate_proxy_recall_url(request_url) 

410 parsed_req = urlparse(request_url) 

411 req_host = parsed_req.hostname 

412 if not req_host: 

413 raise ValueError("proxy recall URL has no host") 

414 req_port = parsed_req.port or (443 if parsed_req.scheme == "https" else 80) 

415 sch = (parsed_req.scheme or "").lower() 

416 

417 pinned = await asyncio.to_thread( 

418 _sync_dns_validate_and_first_public_ip, 

419 req_host, 

420 req_port, 

421 ) 

422 

423 if _is_literal_ip_host(req_host): 

424 effective_url = request_url 

425 out_headers = headers 

426 req_extensions: dict[str, Any] | None = None 

427 else: 

428 effective_url = _rebuild_request_url_pinned_to_ip( 

429 parsed_req, 

430 pinned=pinned, 

431 port=req_port, 

432 ) 

433 out_headers = { 

434 **headers, 

435 "Host": _proxy_recall_host_header_value(req_host, req_port, sch), 

436 } 

437 req_extensions = {"sni_hostname": req_host} if sch == "https" else None 

438 

439 httpx_url, httpx_params = _httpx_url_and_query_params(effective_url) 

440 

441 async with httpx.AsyncClient(timeout=timeout, follow_redirects=False) as client: 

442 if method == "POST": 

443 r = await client.request( 

444 "POST", 

445 httpx_url, 

446 params=httpx_params, 

447 headers=out_headers, 

448 json=body, 

449 extensions=req_extensions, 

450 ) 

451 else: 

452 r = await client.request( 

453 "GET", 

454 httpx_url, 

455 params=httpx_params, 

456 headers=out_headers, 

457 extensions=req_extensions, 

458 ) 

459 r.raise_for_status() 

460 data = r.json() 

461 except Exception: 

462 _record_proxy_metrics( 

463 metrics, source_id=source_id, status="error", duration_s=time.monotonic() - started 

464 ) 

465 raise 

466 duration_s = t["elapsed_ms"] / 1000.0 

467 

468 _record_proxy_metrics(metrics, source_id=source_id, status="ok", duration_s=duration_s) 

469 

470 raw_hits = _parse_hits_payload(data) 

471 out: list[MemoryHit] = [] 

472 for row in raw_hits: 

473 if not isinstance(row, dict): 

474 continue 

475 hit = _row_to_hit(source_id, row) 

476 if hit: 

477 out.append(hit) 

478 return out 

479 

480 

481#: One deadline for all proxy sources of a recall, in seconds. Sources run 

482#: concurrently; whatever has answered by the deadline is fused and the rest 

483#: are cancelled. Before 2026-10-04 sources ran one after another with 15 s 

484#: each, so recall paid the SUM of every remote call (federated-sources §1.1). 

485_DEFAULT_DEADLINE = 0.8 

486#: Consecutive failures (errors or timeouts) before a source is skipped. 

487_BREAKER_THRESHOLD = 3 

488#: How long a tripped source is skipped before it is tried again. 

489_BREAKER_COOLDOWN_S = 60.0 

490 

491# source_id -> (consecutive failures, monotonic time until which it is skipped) 

492_breakers: dict[str, tuple[int, float]] = {} 

493 

494#: A source that misses the deadline keeps running in the background; when it 

495#: answers, its hits are kept here for the next recall of the same query, so 

496#: the work is not thrown away. One-shot: a cached answer is used once. 

497#: Keyed by (source, bank, query): sources authenticate with their own config 

498#: today, not per caller. Per-caller auth (federated-sources F4) must add the 

499#: principal to the key, or one caller would be served another's results. 

500_LATE_TTL_S = 60.0 

501_LATE_MAX_ENTRIES = 256 

502#: A late fetch is cancelled outright after this many deadlines. 

503_LATE_HARD_CAP_FACTOR = 10.0 

504_late_hits: dict[tuple[str, str, str], tuple[float, list[MemoryHit]]] = {} 

505_late_tasks: set[asyncio.Task] = set() 

506 

507 

508def _deadline_seconds() -> float: 

509 raw = os.environ.get("ASTROCYTE_PROXY_RECALL_DEADLINE_SECONDS") 

510 try: 

511 value = float(raw) if raw else _DEFAULT_DEADLINE 

512 except ValueError: 

513 logger.warning("ignoring invalid ASTROCYTE_PROXY_RECALL_DEADLINE_SECONDS=%r", raw) 

514 value = _DEFAULT_DEADLINE 

515 return value if value > 0 else _DEFAULT_DEADLINE 

516 

517 

518def _breaker_open(source_id: str) -> bool: 

519 _failures, until = _breakers.get(source_id, (0, 0.0)) 

520 return time.monotonic() < until 

521 

522 

523def _record_outcome(source_id: str, ok: bool) -> None: 

524 if ok: 

525 _breakers.pop(source_id, None) 

526 return 

527 failures = _breakers.get(source_id, (0, 0.0))[0] + 1 

528 until = time.monotonic() + _BREAKER_COOLDOWN_S if failures >= _BREAKER_THRESHOLD else 0.0 

529 if until: 

530 logger.warning( 

531 "proxy source %s failed %d times in a row; skipping it for %.0f s", 

532 source_id, failures, _BREAKER_COOLDOWN_S, 

533 ) 

534 _breakers[source_id] = (failures, until) 

535 

536 

537def reset_proxy_breakers() -> None: 

538 """Forget every source's failure history and late answers (tests, config reloads).""" 

539 _breakers.clear() 

540 _late_hits.clear() 

541 for task in list(_late_tasks): 

542 task.cancel() 

543 _late_tasks.clear() 

544 

545 

546def _take_late(key: tuple[str, str, str]) -> list[MemoryHit] | None: 

547 entry = _late_hits.pop(key, None) 

548 if entry is None: 

549 return None 

550 expires, hits = entry 

551 return hits if time.monotonic() < expires else None 

552 

553 

554def _keep_late(key: tuple[str, str, str], task: asyncio.Task, hard_cap: asyncio.TimerHandle) -> None: 

555 hard_cap.cancel() 

556 _late_tasks.discard(task) 

557 if task.cancelled() or task.exception() is not None: 

558 return 

559 hits = task.result() 

560 if not hits: 

561 return 

562 _late_hits[key] = (time.monotonic() + _LATE_TTL_S, hits) 

563 while len(_late_hits) > _LATE_MAX_ENTRIES: 

564 _late_hits.pop(next(iter(_late_hits))) 

565 

566 

567async def gather_proxy_hits_for_bank( 

568 config: AstrocyteConfig | Any, 

569 *, 

570 query: str, 

571 bank_id: str, 

572 metrics: MetricsCollector | None = None, 

573) -> list[MemoryHit]: 

574 """Fetch hits from all ``type: proxy`` sources whose ``target_bank`` matches ``bank_id``. 

575 

576 Sources are queried concurrently under one deadline; a source that has not 

577 answered by then contributes nothing to this recall, so one slow or dead 

578 source costs at most the deadline. It keeps running in the background (up 

579 to ten deadlines), and a late answer serves the next recall of the same 

580 query. A source that keeps failing or missing the deadline is skipped for a 

581 cool-down. Hits come back in config order, not completion order, so 

582 fusion ranks do not depend on network timing. 

583 """ 

584 sources = getattr(config, "sources", None) or {} 

585 eligible: list[tuple[str, SourceConfig]] = [] 

586 for sid, src in sources.items(): 

587 if not isinstance(src, SourceConfig): 

588 continue 

589 if (src.type or "").strip().lower() != "proxy": 

590 continue 

591 if (src.target_bank or "").strip() != bank_id: 

592 continue 

593 if _breaker_open(sid): 

594 _record_proxy_metrics(metrics, source_id=sid, status="skipped", duration_s=None) 

595 continue 

596 eligible.append((sid, src)) 

597 if not eligible: 

598 return [] 

599 

600 deadline = _deadline_seconds() 

601 

602 async def one(sid: str, src: SourceConfig) -> list[MemoryHit] | None: 

603 """Hits, or None when the source failed.""" 

604 cap = src.recall_timeout_seconds 

605 timeout = min(cap, deadline) if cap and cap > 0 else deadline 

606 try: 

607 return await fetch_proxy_recall_hits( 

608 sid, src, query=query, bank_id=bank_id, timeout=timeout, metrics=metrics 

609 ) 

610 except Exception as e: 

611 logger.warning("proxy recall failed for source %s: %s", sid, e) 

612 return None 

613 

614 results: dict[str, list[MemoryHit]] = {} 

615 tasks: dict[str, asyncio.Task] = {} 

616 for sid, src in eligible: 

617 late = _take_late((sid, bank_id, query)) 

618 if late is not None: 

619 results[sid] = late 

620 _record_proxy_metrics(metrics, source_id=sid, status="late_cache", duration_s=None) 

621 else: 

622 tasks[sid] = asyncio.create_task(one(sid, src)) 

623 

624 if tasks: 

625 _done, pending = await asyncio.wait(tasks.values(), timeout=deadline) 

626 loop = asyncio.get_running_loop() 

627 for sid, task in tasks.items(): 

628 if task in pending: 

629 # Missing the deadline counts against the source even if it 

630 # answers later: a source that is always late must still trip 

631 # the breaker. Its late answer only fills the cache. 

632 _record_outcome(sid, ok=False) 

633 _record_proxy_metrics(metrics, source_id=sid, status="timeout", duration_s=deadline) 

634 logger.warning("proxy source %s missed the %.2f s recall deadline", sid, deadline) 

635 hard_cap = loop.call_later(deadline * _LATE_HARD_CAP_FACTOR, task.cancel) 

636 _late_tasks.add(task) 

637 task.add_done_callback( 

638 lambda t, k=(sid, bank_id, query), h=hard_cap: _keep_late(k, t, h) 

639 ) 

640 continue 

641 hits = task.result() 

642 _record_outcome(sid, ok=hits is not None) 

643 if hits: 

644 results[sid] = hits 

645 

646 out: list[MemoryHit] = [] 

647 for sid, _src in eligible: 

648 out.extend(results.get(sid, [])) 

649 return out 

650 

651 

652async def merge_manual_and_proxy_hits( 

653 config: AstrocyteConfig | Any, 

654 *, 

655 query: str, 

656 bank_id: str, 

657 manual: list[MemoryHit] | None, 

658 metrics: MetricsCollector | None = None, 

659) -> list[MemoryHit] | None: 

660 """Combine caller ``external_context`` with configured proxy sources for this bank.""" 

661 proxy = await gather_proxy_hits_for_bank(config, query=query, bank_id=bank_id, metrics=metrics) 

662 if not proxy and not manual: 

663 return None 

664 return list(manual or []) + proxy