Skip to content

fix: peer gate race condition + async client connection leak #83

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yasinlex wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: peer gate race condition + async client connection leak #83

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
237 changes: 143 additions & 94 deletions provider_adapters/openai_compatible.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,14 +211,56 @@ def call(request: dict) -> dict:


_PEER_GATES: dict[str, asyncio.Semaphore] = {}


def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore:
sem = _PEER_GATES.get(peer_id)
if sem is None:
sem = asyncio.Semaphore(cap)
_PEER_GATES[peer_id] = sem
return sem
_PEER_GATE_LOCK: asyncio.Lock | None = None

# Maximum idle time (seconds) before a peer gate is evicted to prevent
# unbounded memory growth when peers are dynamically discovered and retired.
_PEER_GATE_TTL: float = 600.0
_PEER_GATE_TIMESTAMPS: dict[str, float] = {}


def _peer_gate_lock() -> asyncio.Lock:
"""Lazy-initialised lock to avoid creating asyncio primitives at import time
(which may happen outside an event loop in test/synchronous contexts)."""
global _PEER_GATE_LOCK
if _PEER_GATE_LOCK is None:
_PEER_GATE_LOCK = asyncio.Lock()
return _PEER_GATE_LOCK


async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore:
"""Return a per-peer concurrency semaphore, creating one if needed.

Uses a lock to prevent a race where two concurrent coroutines both see
``None`` and create separate semaphores for the same peer_id — the
second would overwrite the first, silently dropping any coroutines
already waiting on it.

Old entries are evicted after ``_PEER_GATE_TTL`` seconds of inactivity
so that the dict does not grow without bound when AntSeed peers are
dynamically discovered and later retired.
"""
import time as _time

lock = _peer_gate_lock()
async with lock:
# Evict stale entries first (best-effort, not on every call to
# avoid quadratic sweep cost — only when the lock is already held
# and the dict is getting large).
if len(_PEER_GATES) > 256:
now = _time.monotonic()
stale = [k for k, ts in _PEER_GATE_TIMESTAMPS.items()
if now - ts > _PEER_GATE_TTL]
for k in stale:
_PEER_GATES.pop(k, None)
_PEER_GATE_TIMESTAMPS.pop(k, None)

sem = _PEER_GATES.get(peer_id)
if sem is None:
sem = asyncio.Semaphore(cap)
_PEER_GATES[peer_id] = sem
_PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic()
return sem
Comment on lines 213 to +263

@coderabbitai coderabbitai Bot Jul 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Solid async gate cache design — but cached semaphore ignores a changed cap for the same peer.

The lock-protected creation and TTL eviction are correctly implemented (no await inside the critical section, so no race). However, once a semaphore exists for a peer_id, a later call with a different cap (e.g., seller updates max_concurrency) is silently ignored — sem = _PEER_GATES.get(peer_id) short-circuits before the new cap is ever used, and the stale limit persists until the entry ages out (_PEER_GATE_TTL = 600s).

♻️ Proposed fix: recreate semaphore when cap changes 
+_PEER_GATE_CAPS: dict[str, int] = {}
+
 async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore:
     ...
         sem = _PEER_GATES.get(peer_id)
-        if sem is None:
+        if sem is None or _PEER_GATE_CAPS.get(peer_id) != cap:
             sem = asyncio.Semaphore(cap)
             _PEER_GATES[peer_id] = sem
+            _PEER_GATE_CAPS[peer_id] = cap
         _PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic()
         return sem

Note: in-flight callers still holding the old semaphore continue enforcing the old cap briefly, same as the existing TTL-eviction transition — acceptable tradeoff.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
_PEER_GATES: dict[str, asyncio.Semaphore] = {}
def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore:
sem = _PEER_GATES.get(peer_id)
if sem is None:
sem = asyncio.Semaphore(cap)
_PEER_GATES[peer_id] = sem
return sem
_PEER_GATE_LOCK: asyncio.Lock | None = None
# Maximum idle time (seconds) before a peer gate is evicted to prevent
# unbounded memory growth when peers are dynamically discovered and retired.
_PEER_GATE_TTL: float = 600.0
_PEER_GATE_TIMESTAMPS: dict[str, float] = {}
def _peer_gate_lock() -> asyncio.Lock:
"""Lazy-initialised lock to avoid creating asyncio primitives at import time
(which may happen outside an event loop in test/synchronous contexts)."""
global _PEER_GATE_LOCK
if _PEER_GATE_LOCK is None:
_PEER_GATE_LOCK = asyncio.Lock()
return _PEER_GATE_LOCK
async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore:
"""Return a per-peer concurrency semaphore, creating one if needed.
Uses a lock to prevent a race where two concurrent coroutines both see
``None`` and create separate semaphores for the same peer_id — the
second would overwrite the first, silently dropping any coroutines
already waiting on it.
Old entries are evicted after ``_PEER_GATE_TTL`` seconds of inactivity
so that the dict does not grow without bound when AntSeed peers are
dynamically discovered and later retired.
"""
import time as _time
lock = _peer_gate_lock()
async with lock:
# Evict stale entries first (best-effort, not on every call to
# avoid quadratic sweep cost — only when the lock is already held
# and the dict is getting large).
if len(_PEER_GATES) > 256:
now = _time.monotonic()
stale = [k for k, ts in _PEER_GATE_TIMESTAMPS.items()
if now - ts > _PEER_GATE_TTL]
for k in stale:
_PEER_GATES.pop(k, None)
_PEER_GATE_TIMESTAMPS.pop(k, None)
sem = _PEER_GATES.get(peer_id)
if sem is None:
sem = asyncio.Semaphore(cap)
_PEER_GATES[peer_id] = sem
_PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic()
return sem
_PEER_GATES: dict[str, asyncio.Semaphore] = {}
_PEER_GATE_LOCK: asyncio.Lock | None = None
# Maximum idle time (seconds) before a peer gate is evicted to prevent
# unbounded memory growth when peers are dynamically discovered and retired.
_PEER_GATE_TTL: float = 600.0
_PEER_GATE_TIMESTAMPS: dict[str, float] = {}
_PEER_GATE_CAPS: dict[str, int] = {}
def _peer_gate_lock() -> asyncio.Lock:
"""Lazy-initialised lock to avoid creating asyncio primitives at import time
(which may happen outside an event loop in test/synchronous contexts)."""
global _PEER_GATE_LOCK
if _PEER_GATE_LOCK is None:
_PEER_GATE_LOCK = asyncio.Lock()
return _PEER_GATE_LOCK
async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore:
"""Return a per-peer concurrency semaphore, creating one if needed.
Uses a lock to prevent a race where two concurrent coroutines both see
``None`` and create separate semaphores for the same peer_id — the
second would overwrite the first, silently dropping any coroutines
already waiting on it.
Old entries are evicted after ``_PEER_GATE_TTL`` seconds of inactivity
so that the dict does not grow without bound when AntSeed peers are
dynamically discovered and later retired.
"""
import time as _time
lock = _peer_gate_lock()
async with lock:
# Evict stale entries first (best-effort, not on every call to
# avoid quadratic sweep cost — only when the lock is already held
# and the dict is getting large).
if len(_PEER_GATES) > 256:
now = _time.monotonic()
stale = [k for k, ts in _PEER_GATE_TIMESTAMPS.items()
if now - ts > _PEER_GATE_TTL]
for k in stale:
_PEER_GATES.pop(k, None)
_PEER_GATE_TIMESTAMPS.pop(k, None)
_PEER_GATE_CAPS.pop(k, None)
sem = _PEER_GATES.get(peer_id)
if sem is None or _PEER_GATE_CAPS.get(peer_id) != cap:
sem = asyncio.Semaphore(cap)
_PEER_GATES[peer_id] = sem
_PEER_GATE_CAPS[peer_id] = cap
_PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic()
return sem
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@provider_adapters/openai_compatible.py` around lines 213 - 263, Update
_peer_gate to track the capacity associated with each cached peer semaphore and
recreate the semaphore when the requested cap differs from the cached cap. Keep
the existing semaphore when cap is unchanged, preserve lock protection and
timestamp updates, and ensure replacement entries update the stored capacity so
subsequent calls use the new limit.



def make_async_call_provider(
Expand Down Expand Up @@ -252,7 +294,7 @@ async def call(request: dict) -> dict:
offer = request.get("offer") or {}
peer_id = offer.get("peer_id")
cap = offer.get("max_concurrency")
gate = _peer_gate(peer_id, cap) if (
gate = await _peer_gate(peer_id, cap) if (
peer_id and isinstance(cap, int) and cap > 0) else None
if gate is not None:
try:
Expand Down Expand Up @@ -341,7 +383,8 @@ async def stream_openai_compatible(
body["stream"] = True
rules = (provider_rules or {}).get(request.get("provider_id")) or {}

if client is None:
_owns_client = client is None
if _owns_client:
import httpx
client = httpx.AsyncClient()

Expand All @@ -365,92 +408,98 @@ def _timeout_err() -> dict:
return first_token_timeout_err(first_timeout_s, _latency())

try:
async with AsyncExitStack() as stack:
try:
resp = await before_first_output(stack.enter_async_context(
client.stream("POST", url, json=body, headers=headers,
timeout=timeout)), first_timeout_s, t0, _saw_output)
except (asyncio.TimeoutError, TimeoutError):
return _timeout_err()
if not (200 <= resp.status_code < 300):
raw = (await resp.aread()).decode("utf-8", "replace")[:500]
kind = _classify_from_map(raw, rules.get("error_map")) \
or _classify_status(resp.status_code, raw)
return _err(kind, resp.status_code, _latency(), raw)

lines = resp.aiter_lines().__aiter__()
while True:
try:
async with AsyncExitStack() as stack:
try:
line = await before_first_output(
lines.__anext__(), first_timeout_s, t0, _saw_output)
except StopAsyncIteration:
break
resp = await before_first_output(stack.enter_async_context(
client.stream("POST", url, json=body, headers=headers,
timeout=timeout)), first_timeout_s, t0, _saw_output)
except (asyncio.TimeoutError, TimeoutError):
if not saw_output:
return _timeout_err()
raise
if not line or not line.startswith("data:"):
continue
data = line[len("data:"):].strip()
if data == "[DONE]":
break
try:
chunk = json.loads(data)
except ValueError:
continue
if raw_model is None:
raw_model = chunk.get("model")
if chunk.get("usage"):
usage = chunk["usage"]
for choice in chunk.get("choices") or []:
delta = choice.get("delta") or {}
if choice.get("finish_reason"):
finish_reason = choice["finish_reason"]
content = delta.get("content")
if content:
saw_output = True
text_parts.append(content)
await emit(content)
emitted = True
for tc in delta.get("tool_calls") or []:
saw_output = True
idx = tc.get("index", 0)
acc = tool_calls_acc.setdefault(idx, {
"id": None, "type": "function",
"function": {"name": "", "arguments": ""}})
if tc.get("id"):
acc["id"] = tc["id"]
fn = tc.get("function") or {}
if fn.get("name"):
acc["function"]["name"] = fn["name"]
if fn.get("arguments"):
acc["function"]["arguments"] += fn["arguments"]
except Exception as exc: # noqa: BLE001 — classified below
partial = "".join(text_parts)
if emitted:
return _err("stream_interrupted", 0, _latency(),
f"{type(exc).__name__}: {exc} (partial: {partial[:200]!r})")
return _err("network_error", 0, _latency(), f"{type(exc).__name__}: {exc}")

tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] or None
text = "".join(text_parts)
if not text.strip() and not tool_calls:
return _err("bad_response", 200, _latency(), "empty assistant content")
return {
"ok": True,
"latency_ms": _latency(),
"response": {
"text": text,
"tool_calls": tool_calls,
"finish_reason": finish_reason,
"tokens_in": usage.get("prompt_tokens"),
"tokens_out": usage.get("completion_tokens"),
"tokens_total": usage.get("total_tokens"),
"tokens_cached": _cached_tokens(usage),
"cost_reported": usage.get("cost"),
"raw_model": raw_model,
},
}
return _timeout_err()
if not (200 <= resp.status_code < 300):
raw = (await resp.aread()).decode("utf-8", "replace")[:500]
kind = _classify_from_map(raw, rules.get("error_map")) \
or _classify_status(resp.status_code, raw)
return _err(kind, resp.status_code, _latency(), raw)

lines = resp.aiter_lines().__aiter__()
while True:
try:
line = await before_first_output(
lines.__anext__(), first_timeout_s, t0, _saw_output)
except StopAsyncIteration:
break
except (asyncio.TimeoutError, TimeoutError):
if not saw_output:
return _timeout_err()
raise
if not line or not line.startswith("data:"):
continue
data = line[len("data:"):].strip()
if data == "[DONE]":
break
try:
chunk = json.loads(data)
except ValueError:
continue
if raw_model is None:
raw_model = chunk.get("model")
if chunk.get("usage"):
usage = chunk["usage"]
for choice in chunk.get("choices") or []:
delta = choice.get("delta") or {}
if choice.get("finish_reason"):
finish_reason = choice["finish_reason"]
content = delta.get("content")
if content:
saw_output = True
text_parts.append(content)
await emit(content)
emitted = True
for tc in delta.get("tool_calls") or []:
saw_output = True
idx = tc.get("index", 0)
acc = tool_calls_acc.setdefault(idx, {
"id": None, "type": "function",
"function": {"name": "", "arguments": ""}})
if tc.get("id"):
acc["id"] = tc["id"]
fn = tc.get("function") or {}
if fn.get("name"):
acc["function"]["name"] = fn["name"]
if fn.get("arguments"):
acc["function"]["arguments"] += fn["arguments"]
except Exception as exc: # noqa: BLE001 — classified below
partial = "".join(text_parts)
if emitted:
return _err("stream_interrupted", 0, _latency(),
f"{type(exc).__name__}: {exc} (partial: {partial[:200]!r})")
return _err("network_error", 0, _latency(), f"{type(exc).__name__}: {exc}")

tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] or None
text = "".join(text_parts)
if not text.strip() and not tool_calls:
return _err("bad_response", 200, _latency(), "empty assistant content")
return {
"ok": True,
"latency_ms": _latency(),
"response": {
"text": text,
"tool_calls": tool_calls,
"finish_reason": finish_reason,
"tokens_in": usage.get("prompt_tokens"),
"tokens_out": usage.get("completion_tokens"),
"tokens_total": usage.get("total_tokens"),
"tokens_cached": _cached_tokens(usage),
"cost_reported": usage.get("cost"),
"raw_model": raw_model,
},
}
finally:
# Close the client if we created it, to prevent connection leaks.
# Caller-owned clients are their responsibility.
if _owns_client:
await client.aclose()


def _parse_openai_response(
Expand Down