Spaces:
Running
Running
| """Process-local single-flight coalescer for idempotent upstream calls.""" | |
| from __future__ import annotations | |
| import asyncio | |
| import os | |
| from collections.abc import Awaitable, Callable | |
| from typing import TypeVar | |
| from frontier_agent.infra.usage_meter import record_api_request | |
| T = TypeVar("T") | |
| # Follower wait cap before falling back to an independent call. Bounds the | |
| # "leader stalled → everyone blocks" failure mode. | |
| _DEFAULT_WAIT_TIMEOUT_S = 90.0 | |
| def _enabled() -> bool: | |
| """On by default; set ``TOOL_CALL_COALESCE=0`` (false/off/no) to disable.""" | |
| raw = (os.environ.get("TOOL_CALL_COALESCE") or "").strip().lower() | |
| return raw not in {"0", "false", "off", "no"} | |
| class SingleFlightCoalescer: | |
| """Coalesce concurrent identical in-flight calls into one upstream call. | |
| ``run(key, fn)`` runs ``fn`` once per ``key`` while that call is in flight; | |
| concurrent callers for the same ``key`` (*followers*) await the *leader*'s | |
| result instead of duplicating the call. Nothing is stored after the call | |
| completes — the next caller for the same ``key`` is a fresh leader (and | |
| hits the proxy's own cache). Followers fall back to an independent ``fn`` | |
| call if the leader raises or the wait times out, so coalescing never turns | |
| one transient failure into a correlated N-caller failure. | |
| """ | |
| def __init__(self, name: str = "", *, meter_provider: str = "") -> None: | |
| self.name = name | |
| # Billing provider this coalescer fronts (e.g. | |
| # ``"serper"``). When set, every coalesced follower counts one | |
| # ``external_apis.{provider}.cache_hits`` — an upstream call saved. | |
| self.meter_provider = meter_provider | |
| self._inflight: dict[str, asyncio.Future] = {} | |
| self._lock = asyncio.Lock() | |
| self.leaders = 0 | |
| self.coalesced = 0 | |
| async def run( | |
| self, | |
| key: str, | |
| fn: Callable[[], Awaitable[T]], | |
| *, | |
| wait_timeout: float = _DEFAULT_WAIT_TIMEOUT_S, | |
| ) -> T: | |
| if not key or not _enabled(): | |
| return await fn() | |
| async with self._lock: | |
| fut = self._inflight.get(key) | |
| leader = fut is None | |
| if leader: | |
| fut = asyncio.get_event_loop().create_future() | |
| self._inflight[key] = fut | |
| self.leaders += 1 | |
| else: | |
| self.coalesced += 1 | |
| # A coalesced follower is one upstream call | |
| # saved — surface it as a cache hit on the meter. | |
| if self.meter_provider: | |
| record_api_request( | |
| self.meter_provider, requests=0, cache_hits=1, | |
| ) | |
| if not leader: | |
| try: | |
| # ``shield`` so a follower's timeout can't cancel shared work. | |
| return await asyncio.wait_for( | |
| asyncio.shield(fut), timeout=wait_timeout, | |
| ) | |
| except Exception: | |
| # Leader failed / timed out → independent call (resilience). | |
| return await fn() | |
| # Leader path — run once, hand the result to followers, store nothing. | |
| try: | |
| result = await fn() | |
| except BaseException as exc: | |
| async with self._lock: | |
| self._inflight.pop(key, None) | |
| if not fut.done(): | |
| fut.set_exception(exc) | |
| # Retrieve eagerly so a leader failure with no waiting follower | |
| # doesn't log "Future exception was never retrieved". | |
| fut.exception() | |
| raise | |
| async with self._lock: | |
| self._inflight.pop(key, None) | |
| if not fut.done(): | |
| fut.set_result(result) | |
| return result | |
| def clear(self) -> None: | |
| """Drop in-flight map + counters. Primarily for tests.""" | |
| self._inflight.clear() | |
| self.leaders = self.coalesced = 0 | |
| def stats(self) -> dict[str, int]: | |
| return { | |
| "leaders": self.leaders, | |
| "coalesced": self.coalesced, | |
| "inflight": len(self._inflight), | |
| } | |