Spaces:
Running
Running
File size: 4,112 Bytes
9d0fd45 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """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),
}
|