fix(rate_limiter): use a threading.Lock for the synchronous methods (asyncio.Lock can't use plain with)

[Osamaali313]

Problem

RateLimiter (ms_agent/utils/rate_limiter.py) creates its lock as an asyncio primitive and uses it correctly in the async path:

self._lock = asyncio.Lock()
...
async def _wait_if_needed(self):
    async with self._lock:   # correct

But the synchronous bookkeeping methods acquire that same asyncio lock with a plain with:

• RateLimiter.get_stats / reset
• AdaptiveRateLimiter.record_success / record_error / get_stats

with self._lock:   # asyncio.Lock has no __enter__/__exit__

asyncio.Lock implements only __aenter__/__aexit__, so with self._lock: raises TypeError: 'asyncio.locks.Lock' object does not support the context manager protocol on every call. The correct async with self._lock sibling in the same class shows these were meant to be locked, just not with the async lock.

Impact (live path)

ms_agent/tools/findata/findata_fetcher.py calls record_success() after each successful adaptive-rate-limited fetch and record_error(...) in its except handler. So with rate_limiter: {enabled: true, type: adaptive}, the first successful fetch's record_success() throws TypeError; the fetcher's except then calls record_error(), which throws again and propagates — the fetch fails.

Reproduction

AdaptiveRateLimiter().record_success()

• Before: TypeError: 'asyncio.locks.Lock' object does not support the context manager protocol
• After: succeeds; counters update.

Fix

Add a dedicated threading.Lock for the synchronous methods; the async request-pacing path keeps the asyncio.Lock:

self._lock = asyncio.Lock()
self._state_lock = threading.Lock()
...
# get_stats / reset / record_success / record_error:
with self._state_lock:

AdaptiveRateLimiter.__init__ calls super().__init__, so it inherits _state_lock. Verified RED→GREEN on the extracted logic; python -m py_compile passes. (No existing tests cover rate_limiter; I couldn't import the full class locally because ms_agent.config pulls modelscope, which isn't installed in this env — CI will exercise it.)

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[gemini-code-assist]

Code Review

This pull request introduces a synchronous threading.Lock (self._state_lock) to replace asyncio.Lock in the synchronous bookkeeping methods of RateLimiter (such as get_stats, reset, record_success, and record_error) to prevent runtime errors when using a plain with statement. However, the review highlights two critical issues: first, a synchronization gap exists because the asynchronous _wait_if_needed method still accesses shared state under asyncio.Lock, which can lead to race conditions and runtime errors when synchronous methods are called concurrently from other threads. Second, in AdaptiveRateLimiter, dynamically adjusting max_requests_per_second does not update the immutable maxlen of the underlying deque (self._request_times), which can break the rate-limiting logic.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Synchronization Gap between asyncio.Lock and threading.Lock

While introducing self._state_lock (a threading.Lock) prevents the TypeError in synchronous methods, it introduces a critical thread-safety issue.

The asynchronous _wait_if_needed method still accesses and mutates self._request_times and self._last_request_time under self._lock (an asyncio.Lock), which does not synchronize with self._state_lock.

Since synchronous methods like get_stats, reset, record_success, and record_error can be called from other threads (e.g., when running synchronous tools in a thread pool via run_in_executor), they can access or mutate self._request_times concurrently with _wait_if_needed. This will lead to race conditions and potential RuntimeError: deque mutated during iteration in get_stats.

Recommended Fix

To ensure thread safety without blocking the asyncio event loop during await points, use self._state_lock inside _wait_if_needed to protect all synchronous state accesses and mutations, but release it before any await calls.

Here is how _wait_if_needed should be refactored:

async def _wait_if_needed(self):
    now = time.time()

    # 1. Check minimum request interval safely
    with self._state_lock:
        last_request_time = self._last_request_time
        min_request_interval = self.min_request_interval

    if last_request_time > 0:
        elapsed = now - last_request_time
        if elapsed < min_request_interval:
            wait_time = min_request_interval - elapsed
            logger.debug(f'Enforcing min interval: waiting {wait_time:.3f}s')
            await asyncio.sleep(wait_time)
            now = time.time()

    # 2. Check requests per second limit
    async with self._lock:
        with self._state_lock:
            # Clean up expired request records
            cutoff_time = now - 1.0
            while self._request_times and self._request_times[0] < cutoff_time:
                self._request_times.popleft()

            # Check if rate limit reached
            has_reached_limit = len(self._request_times) >= self.max_requests_per_second
            wait_time = 1.0 - (now - self._request_times[0]) + 0.01 if has_reached_limit else 0

        if has_reached_limit and wait_time > 0:
            logger.debug(f'Rate limit reached: waiting {wait_time:.3f}s')
            await asyncio.sleep(wait_time)
            now = time.time()
            
            with self._state_lock:
                cutoff_time = now - 1.0
                while self._request_times and self._request_times[0] < cutoff_time:
                    self._request_times.popleft()

        with self._state_lock:
            self._request_times.append(now)
            self._last_request_time = now

[gemini-code-assist]

Immutable deque maxlen Bug in AdaptiveRateLimiter

In AdaptiveRateLimiter, the rate limit self.max_requests_per_second is dynamically adjusted in record_success and record_error. However, self._request_times is initialized in the parent class constructor as a deque with a fixed maxlen:

self._request_times = deque(maxlen=max_requests_per_second)

Since a deque's maxlen is immutable after instantiation, changing self.max_requests_per_second does not update the capacity of self._request_times.

• If the rate limit increases (e.g., from 2 to 10), self._request_times remains capped at 2 elements. Consequently, len(self._request_times) will never reach 10, and the rate limiter will never trigger a wait, effectively disabling rate limiting.
• If the rate limit decreases, the deque will still allow more elements than the new limit before discarding.

Recommended Fix

Whenever self.max_requests_per_second is updated, recreate self._request_times with the new maxlen to preserve the existing history up to the new limit:

self._request_times = deque(self._request_times, maxlen=new_rps)

For example, in record_success:

if new_rps > old_rps:
    self.max_requests_per_second = new_rps
    self._request_times = deque(self._request_times, maxlen=new_rps)
    logger.info(...)

And similarly in record_error whenever self.max_requests_per_second is decreased.

[Osamaali313]

Thanks @gemini-code-assist — both are valid, addressed in 18e97f5:

1. Lock synchronization gap. You're right that _state_lock and the asyncio _lock protected the same state independently. _wait_if_needed now guards its _request_times/_last_request_time accesses with _state_lock (the same lock the sync methods use), in brief critical sections that are never held across an await — _lock still serializes coroutines, _state_lock protects the shared state from cross-thread access (and _state_lock is only ever acquired inside _lock, never the reverse, so no lock-ordering cycle).

2. Immutable deque maxlen. Also correct — raising max_requests_per_second never took effect because the deque's maxlen was fixed at the initial rate. Each rate change in record_success/record_error now recreates the deque with the new maxlen, preserving recent history: self._request_times = deque(self._request_times, maxlen=new_rps).

Verified with a standalone async harness (the full suite needs modelscope, which isn't installable in my env): 5 requests at 2 req/s pace to ~2.0s with no deadlock, and the deque capacity now grows from 2→6 on a rate increase (was capped at 2 before).

[gemini-code-assist]

Thanks for the detailed follow-up, @Osamaali313. Your approach to using _state_lock for synchronous bookkeeping while keeping _lock for the async pacing logic is robust, and the explicit handling of the deque recreation ensures that rate limit updates are correctly applied. This separation of concerns correctly avoids potential deadlocks while ensuring thread safety for the shared state. Great work on the fix.