Preserve token invalidation across concurrent refreshes

src/anthropic/lib/credentials/_cache.py

@@ -69,6 +69,10 @@ def __init__(
         # One-shot: invalidate() sets it; next provider call passes
         # force_refresh=True so on-disk providers don't re-serve a stale token.
         self._next_force = False
+        # Monotonic invalidation generation. A refresh snapshots this before it
+        # leaves the lock; if invalidate() advances it while the provider call
+        # is in flight, that result must not be published or returned.
+        self._invalidation_generation = 0
         # Time of last advisory-refresh failure (never reset on success —
         # only distance-from-now matters).
         self._last_advisory_failure_time: float = 0.0
@@ -85,28 +89,23 @@ def _invoke_provider(self, *, force: bool) -> AccessToken:
                 raise
             return self._provider()  # type: ignore[call-arg]
 
-    def _call_provider(self) -> AccessToken:
+    def _call_provider(self, *, force: bool) -> AccessToken:
         """Call the provider, retrying once on a 401 from the token endpoint."""
-        # Read but don't clear yet — clearing only on success keeps the flag
-        # alive across a transient failure so the retry still forces.
-        with self._lock:
-            force = self._next_force
         try:
-            result = self._invoke_provider(force=force)
+            return self._invoke_provider(force=force)
         except WorkloadIdentityError as err:
             if err.status_code != 401:
                 raise
             log.debug("Token provider returned 401; retrying once")
-            result = self._invoke_provider(force=True)
-        with self._lock:
-            self._next_force = False
-        return result
+            return self._invoke_provider(force=True)
 
     def get_token(self) -> str:
         """Return a valid bearer token, refreshing if necessary."""
         while True:
             advisory_fallback: Optional[AccessToken] = None
             remaining_seconds = 0
+            refresh_generation: Optional[int] = None
+            force_refresh = False
             with self._lock:
                 cached = self._cached
                 if cached is not None:
@@ -132,9 +131,13 @@ def get_token(self) -> str:
                     # Mandatory-window caller with a refresh in flight: wait.
                     waiter_event: Optional[threading.Event] = self._refresh_event
                 else:
-                    # We're the leader.
+                    # We're the leader. Snapshot both invalidation state and the
+                    # one-shot force flag before leaving the lock for the
+                    # provider call.
                     self._refresh_event = threading.Event()
                     waiter_event = None
+                    refresh_generation = self._invalidation_generation
+                    force_refresh = self._next_force
 
             if waiter_event is not None:
                 waiter_event.wait()
@@ -143,19 +146,26 @@ def get_token(self) -> str:
                 # refresh ourselves), or been invalidated in between.
                 continue
 
+            assert refresh_generation is not None
+
             # Leader: run the provider outside the lock. The except catches
             # BaseException (not a narrow tuple) so the refresh event is
             # always released — a user-supplied provider raising e.g.
             # RuntimeError must not deadlock mandatory-window waiters.
             try:
-                fresh = self._call_provider()
+                fresh = self._call_provider(force=force_refresh)
             except BaseException as err:
                 with self._lock:
+                    invalidated = self._invalidation_generation != refresh_generation
                     released = self._refresh_event
                     self._refresh_event = None
                 assert released is not None
                 released.set()
-                if advisory_fallback is not None and isinstance(err, (AnthropicError, httpx.HTTPError)):
+                if (
+                    not invalidated
+                    and advisory_fallback is not None
+                    and isinstance(err, (AnthropicError, httpx.HTTPError))
+                ):
                     log.warning(
                         "Advisory token refresh failed (%ds remaining); serving cached token: %s",
                         remaining_seconds,
@@ -167,19 +177,40 @@ def get_token(self) -> str:
                 raise
 
             with self._lock:
-                self._cached = fresh
+                invalidated = self._invalidation_generation != refresh_generation
+                if not invalidated:
+                    self._cached = fresh
+                    # Consume the one-shot force only if this refresh still
+                    # belongs to the current invalidation generation.
+                    self._next_force = False
                 released = self._refresh_event
                 self._refresh_event = None
             assert released is not None
             released.set()
+            if invalidated:
+                log.debug("Discarding token refresh result invalidated while provider call was in flight")
+                # invalidate() cleared the cache and left _next_force set. Loop
+                # so this caller participates in the next single-flight forced
+                # refresh instead of returning the superseded token.
+                continue
             return fresh.token
 
     def invalidate(self) -> None:
         """Clear the cached token so the next :meth:`get_token` re-invokes the provider.
 
         Also sets a one-shot ``force_refresh`` flag so on-disk providers skip
         their freshness short-circuit instead of re-serving the revoked token.
+        A refresh already in flight is invalidated as well: its result is
+        discarded when it returns and cannot consume the force flag.
+
+        Repeated invalidations are coalesced while the cache is already empty
+        and a forced refresh is pending. This prevents several requests that
+        all receive 401s for the same revoked token from invalidating the
+        replacement refresh over and over.
         """
         with self._lock:
+            already_invalidated = self._cached is None and self._next_force
             self._cached = None
             self._next_force = True
+            if not already_invalidated:
+                self._invalidation_generation += 1

tests/lib/test_token_cache_invalidation.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import threading
+
+from anthropic import AccessToken, AnthropicError, TokenCache
+
+
+def test_invalidate_during_refresh_discards_result_and_preserves_singleflight() -> None:
+    first_started = threading.Event()
+    release_first = threading.Event()
+    forced_started = threading.Event()
+    release_forced = threading.Event()
+    calls_lock = threading.Lock()
+    force_seen: list[bool] = []
+
+    def provider(*, force_refresh: bool = False) -> AccessToken:
+        with calls_lock:
+            index = len(force_seen)
+            force_seen.append(force_refresh)
+
+        if index == 0:
+            first_started.set()
+            assert release_first.wait(timeout=5)
+            return AccessToken("pre-invalidation", expires_at=None)
+        if index == 1:
+            forced_started.set()
+            assert release_forced.wait(timeout=5)
+            return AccessToken("post-invalidation", expires_at=None)
+        raise AssertionError(f"unexpected provider call {index + 1}")
+
+    cache = TokenCache(provider)
+    results: list[str] = []
+    errors: list[BaseException] = []
+
+    def get_token() -> None:
+        try:
+            results.append(cache.get_token())
+        except BaseException as exc:  # pragma: no cover - assertion below reports the failure
+            errors.append(exc)
+
+    first = threading.Thread(target=get_token, daemon=True)
+    second = threading.Thread(target=get_token, daemon=True)
+    first.start()
+    assert first_started.wait(timeout=5)
+
+    # The second caller must join the same in-flight refresh rather than start a
+    # parallel provider call. Invalidate while both callers depend on that
+    # refresh, then let its pre-invalidation result return.
+    second.start()
+    cache.invalidate()
+    release_first.set()
+
+    # The stale result must be discarded. Exactly one caller becomes the next
+    # single-flight leader and performs a forced refresh; the other waits.
+    assert forced_started.wait(timeout=5)
+
+    # Another request may now report a 401 for the same revoked token. That
+    # duplicate invalidation is already represented by the pending forced
+    # refresh and must not invalidate the replacement refresh itself.
+    cache.invalidate()
+    release_forced.set()
+
+    first.join(timeout=5)
+    second.join(timeout=5)
+    assert not first.is_alive()
+    assert not second.is_alive()
+    assert errors == []
+    assert sorted(results) == ["post-invalidation", "post-invalidation"]
+    assert force_seen == [False, True]
+
+    # The forced result was published as the cache value; no third provider
+    # call is needed.
+    assert cache.get_token() == "post-invalidation"
+    assert force_seen == [False, True]
+
+
+def test_invalidate_during_failed_advisory_refresh_does_not_serve_stale_fallback() -> None:
+    refresh_started = threading.Event()
+    release_refresh = threading.Event()
+    calls_lock = threading.Lock()
+    force_seen: list[bool] = []
+
+    def provider(*, force_refresh: bool = False) -> AccessToken:
+        with calls_lock:
+            index = len(force_seen)
+            force_seen.append(force_refresh)
+
+        if index == 0:
+            # At t=100 this sits in the advisory window: 50s remaining, with
+            # mandatory=10 and advisory=100.
+            return AccessToken("cached", expires_at=150)
+        if index == 1:
+            refresh_started.set()
+            assert release_refresh.wait(timeout=5)
+            raise AnthropicError("refresh failed")
+        if index == 2:
+            return AccessToken("forced-fresh", expires_at=None)
+        raise AssertionError(f"unexpected provider call {index + 1}")
+
+    cache = TokenCache(
+        provider,
+        advisory_refresh_seconds=100,
+        mandatory_refresh_seconds=10,
+        time_source=lambda: 100,
+    )
+    assert cache.get_token() == "cached"
+
+    results: list[str] = []
+    errors: list[BaseException] = []
+
+    def advisory_refresh() -> None:
+        try:
+            results.append(cache.get_token())
+        except BaseException as exc:
+            errors.append(exc)
+
+    thread = threading.Thread(target=advisory_refresh, daemon=True)
+    thread.start()
+    assert refresh_started.wait(timeout=5)
+
+    # The cached token is revoked while the advisory refresh is in flight. A
+    # subsequent provider failure must not use that revoked token as the normal
+    # advisory fallback.
+    cache.invalidate()
+    release_refresh.set()
+    thread.join(timeout=5)
+
+    assert not thread.is_alive()
+    assert results == []
+    assert len(errors) == 1
+    assert isinstance(errors[0], AnthropicError)
+
+    # The invalidation's one-shot force flag survives the failed in-flight
+    # refresh and is consumed by the next successful provider call.
+    assert cache.get_token() == "forced-fresh"
+    assert force_seen == [False, False, True]