RomanNum3ral
/
hermes-agent
mirror of https://github.com/NousResearch/hermes-agent.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): keep pending turn lease acquires registered

This commit is contained in:
poisdahl 2026-08-06 22:33:48 +02:00 committed by kshitij
parent b3e9e91709
commit 3a3aed3c1f
2 changed files with 230 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
gateway/turn_lease.py
View File

@ -55,10 +55,10 @@ from typing import Dict, Optional
logger = logging.getLogger(__name__)
# Upper bound on tracked per-session leases. Idle entries (no holder, no
# waiter) are evicted oldest-first once the cap is reached; live leases are
# never evicted, so a burst of distinct sessions can transiently exceed the
# cap rather than break serialization.
# Upper bound on tracked per-session leases. Idle entries (no holder or
# pending acquire) are evicted oldest-first once the cap is reached; live
# leases are never evicted, so a burst of distinct sessions can transiently
# exceed the cap rather than break serialization.
DEFAULT_MAX_LEASES = 512
# Fallback wait (seconds) when the caller passes no positive timeout. The
@ -127,18 +127,29 @@ class TurnLeaseToken:
class _SessionLease:
__slots__ = ("lock", "holder", "acquired_at", "last_used")
__slots__ = (
"lock",
"holder",
"acquired_at",
"last_used",
"pending_acquires",
)
def __init__(self) -> None:
self.lock = asyncio.Lock()
self.holder: Optional[TurnLeaseToken] = None
self.acquired_at = 0.0
self.last_used = time.time()
self.pending_acquires = 0
@property
def idle(self) -> bool:
"""True when this lease can be evicted: nobody holds or awaits it."""
return self.holder is None and not self.lock.locked()
return (
self.holder is None
and not self.lock.locked()
and self.pending_acquires == 0
)
class SessionTurnLeaseRegistry:
@ -217,6 +228,12 @@ class SessionTurnLeaseRegistry:
time.time() - lease.acquired_at if lease.acquired_at else -1.0,
)
# Lock.release() wakes a waiter while leaving the lock momentarily
# unlocked. Track every in-progress acquire across that handoff so
# eviction cannot orphan the old lock and create a second lock for the
# same session. Count even apparently-uncontended acquires: wait_for()
# may schedule them before the underlying lock coroutine runs.
lease.pending_acquires += 1
try:
await asyncio.wait_for(lease.lock.acquire(), timeout=wait)
except asyncio.TimeoutError:
@ -239,7 +256,11 @@ class SessionTurnLeaseRegistry:
generation=generation,
wait_seconds=wait,
) from None
finally:
lease.pending_acquires -= 1
# The lock is held and there is no await before holder publication, so
# the lease cannot become evictable after the pending count is cleared.
lease.holder = token
lease.acquired_at = time.time()
lease.last_used = lease.acquired_at

204
tests/gateway/test_turn_lease.py
View File

@ -14,7 +14,8 @@ Covers:
- timeout fail-closed: a timed-out waiter never enters the transcript region,
and outer dispatch returns a visible rejection/resend notice without invoking
goal continuation
- registry stays bounded; live leases are never evicted
- registry stays bounded; live and pending leases are never evicted
- timed-out and cancelled acquire attempts do not pin idle registry entries
- GatewayRunner._release_turn_lease wiring (bare-runner safe, token-scoped)
"""
@ -216,6 +217,172 @@ async def test_full_dispatch_rejects_lease_timeout_without_running_goal_hook(
# ---------------------------------------------------------------------------
class _ObservedLock:
"""Forwarding lock that signals when an acquire has to wait."""
def __init__(self, lock):
self._lock = lock
self.blocked = asyncio.Queue()
async def acquire(self):
if self._lock.locked():
self.blocked.put_nowait(None)
return await self._lock.acquire()
def locked(self):
return self._lock.locked()
def release(self):
self._lock.release()
class _GatedLock:
"""Forwarding lock that pauses an otherwise-uncontended acquire."""
def __init__(self, lock):
self._lock = lock
self.started = asyncio.Event()
self.proceed = asyncio.Event()
async def acquire(self):
self.started.set()
await self.proceed.wait()
return await self._lock.acquire()
def locked(self):
return self._lock.locked()
def release(self):
self._lock.release()
def test_registry_does_not_evict_lease_during_waiter_handoff():
"""A woken waiter must stay in the original serialization domain.
``asyncio.Lock.release()`` unlocks before the selected waiter resumes.
Capacity eviction in that handoff window must not orphan the old lock and
let a later acquire for the same session take a second lock concurrently.
"""
async def scenario():
registry = SessionTurnLeaseRegistry(max_entries=1)
first = await registry.acquire(
"shared", owner_key="first", generation=1, timeout=1
)
lease = registry._leases["shared"]
observed = _ObservedLock(lease.lock)
lease.lock = observed
waking_task = asyncio.create_task(
registry.acquire("shared", owner_key="waking", generation=1, timeout=1)
)
await asyncio.wait_for(observed.blocked.get(), timeout=1)
assert registry.release(first) is True
# _get_or_create("other") runs before this acquire first yields, in
# the unlocked handoff window before waking_task resumes.
other = await registry.acquire(
"other", owner_key="other", generation=1, timeout=1
)
waking = await waking_task
assert registry._leases.get("shared") is lease
successor_task = asyncio.create_task(
registry.acquire("shared", owner_key="successor", generation=2, timeout=1)
)
await asyncio.wait_for(observed.blocked.get(), timeout=1)
assert not successor_task.done()
assert registry.release(waking) is True
successor = await successor_task
assert registry.release(successor) is True
assert registry.release(other) is True
_run(scenario())
def test_registry_does_not_evict_an_uncontended_acquire_before_it_locks():
"""Every pending acquire is protected, even if the lock looked free."""
async def scenario():
registry = SessionTurnLeaseRegistry(max_entries=1)
seed = await registry.acquire(
"shared", owner_key="seed", generation=1, timeout=1
)
assert registry.release(seed) is True
lease = registry._leases["shared"]
gated = _GatedLock(lease.lock)
lease.lock = gated
pending_task = asyncio.create_task(
registry.acquire("shared", owner_key="pending", generation=2, timeout=1)
)
await asyncio.wait_for(gated.started.wait(), timeout=1)
other = await registry.acquire(
"other", owner_key="other", generation=1, timeout=1
)
assert registry._leases.get("shared") is lease
gated.proceed.set()
pending = await pending_task
assert registry.release(pending) is True
assert registry.release(other) is True
_run(scenario())
def test_timed_out_acquire_does_not_pin_idle_registry_entry():
async def scenario():
registry = SessionTurnLeaseRegistry(max_entries=1)
holder = await registry.acquire(
"shared", owner_key="holder", generation=1, timeout=1
)
with pytest.raises(TurnLeaseTimeoutError):
await registry.acquire(
"shared", owner_key="timeout", generation=2, timeout=0.02
)
assert registry.release(holder) is True
other = await registry.acquire(
"other", owner_key="other", generation=1, timeout=1
)
assert set(registry._leases) == {"other"}
assert registry.release(other) is True
_run(scenario())
def test_cancelled_acquire_does_not_pin_idle_registry_entry():
async def scenario():
registry = SessionTurnLeaseRegistry(max_entries=1)
holder = await registry.acquire(
"shared", owner_key="holder", generation=1, timeout=1
)
lease = registry._leases["shared"]
observed = _ObservedLock(lease.lock)
lease.lock = observed
cancelled_task = asyncio.create_task(
registry.acquire("shared", owner_key="cancelled", generation=2, timeout=1)
)
await asyncio.wait_for(observed.blocked.get(), timeout=1)
cancelled_task.cancel()
with pytest.raises(asyncio.CancelledError):
await cancelled_task
assert registry.release(holder) is True
other = await registry.acquire(
"other", owner_key="other", generation=1, timeout=1
)
assert set(registry._leases) == {"other"}
assert registry.release(other) is True
_run(scenario())
# ---------------------------------------------------------------------------
# Mid-turn rotation rebind
# ---------------------------------------------------------------------------
@ -248,6 +415,41 @@ def test_rebind_moves_serialization_to_new_session_id():
_run(scenario())
def test_rebind_does_not_replace_target_during_waiter_handoff():
"""A target with a waking waiter is still a live lease domain."""
async def scenario():
registry = SessionTurnLeaseRegistry()
target_holder = await registry.acquire(
"target", owner_key="target-holder", generation=1, timeout=1
)
target_lease = registry._leases["target"]
observed = _ObservedLock(target_lease.lock)
target_lease.lock = observed
target_waiter_task = asyncio.create_task(
registry.acquire(
"target", owner_key="target-waiter", generation=2, timeout=1
)
)
await asyncio.wait_for(observed.blocked.get(), timeout=1)
source_holder = await registry.acquire(
"source", owner_key="source-holder", generation=1, timeout=1
)
assert registry.release(target_holder) is True
# The target lock is briefly unlocked, but its selected waiter has
# not resumed. Rebind must not replace that serialization domain.
assert registry.rebind(source_holder, "target") is False
target_waiter = await target_waiter_task
assert registry._leases["target"] is target_lease
assert registry.release(target_waiter) is True
assert registry.release(source_holder) is True
_run(scenario())
# ---------------------------------------------------------------------------
# GatewayRunner wiring
# ---------------------------------------------------------------------------