fix(providers): claim probe workers under the pool lock (#3787)

api/providers.py

@@ -1354,6 +1354,7 @@ class _AccountUsageProbeWorker:
                 env=_account_usage_subprocess_env(self.home, provider, None),
                 **kwargs,
             )
+            self._closed = False
         except Exception:
             self._proc = None
             logger.debug("Account usage worker for %s failed to launch", provider, exc_info=True)
@@ -1434,15 +1435,26 @@ def _get_account_usage_probe_worker(home: Path) -> "_AccountUsageProbeWorker | N
     let two concurrent probes both observe the same worker as free.
     """
     key = str(Path(home))
+    stale: list[_AccountUsageProbeWorker] = []
+    claimed: _AccountUsageProbeWorker | None = None
     with _account_usage_worker_pool_lock:
-        workers = _account_usage_worker_pool.get(key)
+        existing_workers = _account_usage_worker_pool.get(key)
-        if not workers:
+        workers: list[_AccountUsageProbeWorker]
+        if not existing_workers:
             workers = [_AccountUsageProbeWorker(Path(home)) for _ in range(_ACCOUNT_USAGE_WORKERS_PER_HOME)]
-            _account_usage_worker_pool[key] = workers
+        else:
-    for w in workers:
+            stale = [worker for worker in existing_workers if worker._closed]
-        if w._lock.acquire(blocking=False):
+            workers = [worker for worker in existing_workers if not worker._closed]
-            return w
+            while len(workers) < _ACCOUNT_USAGE_WORKERS_PER_HOME:
-    return None
+                workers.append(_AccountUsageProbeWorker(Path(home)))
+        _account_usage_worker_pool[key] = workers
+        for worker in workers:
+            if worker._lock.acquire(blocking=False):
+                claimed = worker
+                break
+    for worker in stale:
+        worker.close()
+    return claimed
 
 
 def _cleanup_account_usage_probe_workers(

tests/test_issue3787_probe_pool.py

@@ -323,6 +323,92 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase):
             release_signal.set()
             holder_thread.join(timeout=1.0)
 
+    def test_invalidation_cannot_pop_worker_between_lookup_and_claim(self):
+        """Getter keeps the pool lock through worker claim, so invalidation cannot pop first."""
+        home = Path("/tmp/test_claim_under_pool_lock")
+
+        initial = _get_account_usage_probe_worker(home)
+        self.assertIsNotNone(initial)
+        initial._lock.release()
+
+        with _account_usage_worker_pool_lock:
+            workers = _account_usage_worker_pool[str(Path(home))]
+
+        workers[1]._lock.acquire()
+        self.addCleanup(workers[1]._lock.release)
+
+        acquire_entered = threading.Event()
+        allow_acquire = threading.Event()
+        claimed = threading.Event()
+        release_claim = threading.Event()
+        invalidation_done = threading.Event()
+        result: dict[str, object] = {}
+        target = workers[0]
+        original_lock = target._lock
+        self_outer = self
+
+        class GateLock:
+            def __init__(self, inner):
+                self._inner = inner
+
+            def acquire(self, blocking=True, timeout=-1):
+                if blocking is False:
+                    self_outer.assertTrue(
+                        _account_usage_worker_pool_lock.locked(),
+                        "worker claim must happen while the pool lock is still held",
+                    )
+                    acquire_entered.set()
+                    allow_acquire.wait(timeout=2.0)
+                if timeout == -1:
+                    return self._inner.acquire(blocking)
+                return self._inner.acquire(blocking, timeout)
+
+            def release(self):
+                return self._inner.release()
+
+            def __enter__(self):
+                self._inner.acquire()
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                self._inner.release()
+                return False
+
+        target._lock = GateLock(original_lock)
+
+        def run_getter():
+            result["worker"] = _get_account_usage_probe_worker(home)
+            claimed.set()
+            release_claim.wait(timeout=2.0)
+            if result["worker"] is not None:
+                result["worker"]._lock.release()
+
+        def run_invalidation():
+            with mock.patch("api.providers._get_hermes_home", return_value=home):
+                invalidate_account_usage_status_cache(provider_id="anthropic")
+            invalidation_done.set()
+
+        fetch_thread = threading.Thread(target=run_getter, daemon=True)
+        invalidation_thread = threading.Thread(target=run_invalidation, daemon=True)
+        fetch_thread.start()
+        self.assertTrue(acquire_entered.wait(timeout=2.0))
+
+        invalidation_thread.start()
+        time.sleep(0.05)
+        self.assertFalse(invalidation_done.is_set())
+
+        allow_acquire.set()
+        self.assertTrue(claimed.wait(timeout=2.0))
+        self.assertIs(result["worker"], target)
+
+        release_claim.set()
+        fetch_thread.join(timeout=2.0)
+
+        invalidation_thread.join(timeout=2.0)
+        self.assertTrue(invalidation_done.is_set())
+        with _account_usage_worker_pool_lock:
+            self.assertNotIn(str(Path(home)), _account_usage_worker_pool)
+
 
 class TestWorkerPoolConfiguration(unittest.TestCase):
     """Test that constants are properly configured."""