diff --git a/backend/services/agent_shell_ws_token.py b/backend/services/agent_shell_ws_token.py index 9653b98..72ec452 100644 --- a/backend/services/agent_shell_ws_token.py +++ b/backend/services/agent_shell_ws_token.py @@ -13,25 +13,31 @@ _store: dict[str, float] = {} _lock = Lock() -def _purge_expired(*, force: bool = False) -> None: +def _purge_expired_locked(now: float) -> None: + """Drop expired tokens while the caller holds ``_lock``.""" + expired = [token for token, expires in _store.items() if expires <= now] + for token in expired: + _store.pop(token, None) + + +def _purge_expired() -> None: now = time.time() with _lock: - expired = [token for token, expires in _store.items() if expires <= now] - for token in expired: - _store.pop(token, None) - if force and len(_store) > _MAX_ACTIVE_TOKENS: - for token in list(_store.keys())[: len(_store) - _MAX_ACTIVE_TOKENS]: - _store.pop(token, None) + _purge_expired_locked(now) def mint_agent_shell_ws_token() -> tuple[str, int]: - """Return (token, expires_in_seconds).""" - _purge_expired() + """Return (token, expires_in_seconds) without exceeding the store bound.""" token = secrets.token_urlsafe(32) - expires_at = time.time() + _TOKEN_TTL_SECONDS + now = time.time() + expires_at = now + _TOKEN_TTL_SECONDS with _lock: + _purge_expired_locked(now) if len(_store) >= _MAX_ACTIVE_TOKENS: - _purge_expired(force=True) + # Evict the token that will expire soonest. This keeps the store + # bounded without recursively acquiring the non-reentrant lock. + oldest = min(_store, key=_store.get) + _store.pop(oldest, None) _store[token] = expires_at return token, int(_TOKEN_TTL_SECONDS) @@ -43,6 +49,7 @@ def consume_agent_shell_ws_token(token: str) -> bool: return False now = time.time() with _lock: + _purge_expired_locked(now) expires_at = _store.pop(cleaned, None) return expires_at is not None and expires_at > now