feat(core): add KeyAuth.close() to stop the token refresh thread (#10303)

relates to #10253

---------

Co-authored-by: cgoetz-inovex <[email protected]>

Author: mittalpk

CHANGELOG.md

@@ -1,6 +1,8 @@
 ## Release (2026-MM-DD)
 
 - `core`:
+  - [v0.4.0](core/CHANGELOG.md#v040)
+    - **Feature:** Add `KeyAuth.close()` to stop the background token refresh thread and its worker, so a discarded instance no longer keeps running and requesting tokens
   - [v0.3.1](core/CHANGELOG.md#v031)
     - **Fix:** Refresh access tokens five minutes before expiration instead of after, and wait for in-progress refreshes before authorizing requests
 - `alb`:

core/CHANGELOG.md

@@ -1,3 +1,6 @@
+## v0.4.0
+- **Feature:** Add `KeyAuth.close()` to stop the background token refresh thread and its worker, so a discarded instance (for example after key rotation) no longer keeps running and requesting tokens
+
 ## v0.3.1
 - **Fix:** Refresh access tokens five minutes before expiration instead of after, and wait for in-progress refreshes before authorizing requests
 

core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stackit-core"
-version = "v0.3.1"
+version = "v0.4.0"
 description = "Core functionality for the STACKIT SDK for Python"
 authors = [{ name = "STACKIT Developer Tools", email = "[email protected]" }]
 requires-python = ">=3.10,<4.0"

core/src/stackit/core/auth_methods/key_auth.py

@@ -54,6 +54,8 @@ class KeyAuth(AuthBase):
     executor: ThreadPoolExecutor
     refresh_future: Optional[futures.Future]
     service_account_key: ServiceAccountKey
+    _stop: threading.Event
+    _refresh_thread: Optional[threading.Thread]
 
     def __init__(
         self,
@@ -66,6 +68,8 @@ def __init__(
         self.executor = ThreadPoolExecutor(max_workers=1)
         self.access_token = None
         self.refresh_future = None
+        self._stop = threading.Event()
+        self._refresh_thread = None
         self.__create_initial_token()
         self.__fetch_token_from_endpoint()
         self.__start_token_refresh_task()
@@ -92,6 +96,24 @@ def __call__(self, r: Request) -> Request:
             r.headers["Authorization"] = f"Bearer {self.access_token}"
         return r
 
+    def close(self) -> None:
+        """Stop the background token refresh and release its worker.
+
+        Call this when the instance is no longer needed, for example after
+        replacing it during key rotation in a long-running process. Without it
+        the refresh thread keeps running and keeps requesting tokens. When it
+        returns, no refresh is running and none will start. Safe to call more
+        than once.
+        """
+        # Set under the lock so the refresh loop cannot submit after shutdown.
+        with self.lock:
+            self._stop.set()
+        atexit.unregister(self.__shutdown)
+        thread = self._refresh_thread
+        if thread is not None and thread is not threading.current_thread():
+            thread.join()
+        self.executor.shutdown(wait=True, cancel_futures=True)
+
     def __create_initial_token(self) -> None:
         payload = {
             "iss": self.service_account_key.credentials.issuer,
@@ -126,9 +148,10 @@ def __fetch_token_from_endpoint(self) -> None:
 
     def __start_token_refresh_task(self):
         def token_refresh_task():
-            while True:
-                time.sleep(self.TOKEN_EXPIRY_CHECK_INTERVAL.total_seconds())
+            while not self._stop.wait(self.TOKEN_EXPIRY_CHECK_INTERVAL.total_seconds()):
                 with self.lock:
+                    if self._stop.is_set():
+                        return
                     if self.refresh_future is not None and self.refresh_future.done():
                         self.refresh_future.result()
                     if self.__is_token_expired(self.access_token) and (
@@ -139,6 +162,7 @@ def token_refresh_task():
         thread = threading.Thread(target=token_refresh_task)
         thread.daemon = True
         thread.start()
+        self._refresh_thread = thread
 
     def __refresh_token(self):
         with self.lock:
@@ -156,6 +180,8 @@ def __refresh_token(self):
 
         last_exception = None
         for _ in range(self.MAX_REFRESH_RETRIES):
+            if self._stop.is_set():
+                return
             try:
                 response = requests.post(self.token_endpoint, data=body, timeout=self.timeout)
                 response.raise_for_status()

core/tests/core/test_auth.py

@@ -2,6 +2,9 @@
 from pathlib import Path, PurePath
 from threading import Event, Lock, Thread
 
+import time
+from datetime import timedelta
+
 import pytest
 import json
 import jwt
@@ -308,6 +311,7 @@ def test_call_propagates_token_refresh_failure(self, refresh_already_failed):
         auth.refresh_token = jwt.encode({"exp": 4_000_000_000}, "x" * 32, algorithm="HS256")
         auth.token_endpoint = KeyAuth.DEFAULT_TOKEN_ENDPOINT
         auth.refresh_future = None
+        auth._stop = Event()
         original_token = auth.access_token
         request = Request("GET", "https://example.com")
         refresh_error = requests.RequestException("refresh failed")
@@ -382,3 +386,90 @@ def result(self):
 
         assert not call_thread.is_alive()
         assert request.headers["Authorization"] == f"Bearer {fresh_token}"
+
+
+class TestKeyAuthClose:
+    """A discarded KeyAuth must be able to stop its background token refresh."""
+
+    @staticmethod
+    def _make_auth(service_account_key_file_json, mock_post, access_token):
+        service_account_key = ServiceAccountKey.model_validate_json(service_account_key_file_json)
+        service_account_key.credentials.private_key = "test-private-key"
+        response = Mock()
+        response.raise_for_status.return_value = None
+        response.json.return_value = {
+            "access_token": access_token,
+            "refresh_token": jwt.encode({"exp": 4102444800}, "secret", algorithm="HS256"),
+        }
+        mock_post.return_value = response
+        return KeyAuth(service_account_key)
+
+    def test_close_stops_the_refresh_thread_and_further_token_requests(self, service_account_key_file_json):
+        def set_initial_token(auth):
+            auth.initial_token = "test-initial-token"
+
+        expired = jwt.encode({"exp": 1}, "secret", algorithm="HS256")
+        with (
+            patch.object(KeyAuth, "_KeyAuth__create_initial_token", new=set_initial_token),
+            patch.object(KeyAuth, "TOKEN_EXPIRY_CHECK_INTERVAL", timedelta(milliseconds=10)),
+            patch("requests.post") as mock_post,
+        ):
+            # An access token that is always expired keeps the refresh loop busy.
+            auth = self._make_auth(service_account_key_file_json, mock_post, expired)
+            thread = auth._refresh_thread
+            assert thread is not None and thread.is_alive()
+
+            auth.close()
+
+            assert not thread.is_alive()
+            requests_after_close = mock_post.call_count
+            time.sleep(0.1)  # ten refresh intervals
+            assert mock_post.call_count == requests_after_close
+
+    def test_close_waits_for_a_refresh_in_progress(self, service_account_key_file_json):
+        def set_initial_token(auth):
+            auth.initial_token = "test-initial-token"
+
+        expired = jwt.encode({"exp": 1}, "secret", algorithm="HS256")
+        refresh_started = Event()
+        with (
+            patch.object(KeyAuth, "_KeyAuth__create_initial_token", new=set_initial_token),
+            patch.object(KeyAuth, "TOKEN_EXPIRY_CHECK_INTERVAL", timedelta(milliseconds=10)),
+            patch("requests.post") as mock_post,
+        ):
+            auth = self._make_auth(service_account_key_file_json, mock_post, expired)
+            initial_response = mock_post.return_value
+
+            def slow_refresh(*args, **kwargs):
+                refresh_started.set()
+                time.sleep(0.2)
+                return initial_response
+
+            mock_post.side_effect = slow_refresh
+            assert refresh_started.wait(timeout=5)
+
+            auth.close()
+
+            assert auth.refresh_future is not None and auth.refresh_future.done()
+            requests_after_close = mock_post.call_count
+            time.sleep(0.1)
+            assert mock_post.call_count == requests_after_close
+
+    def test_close_is_idempotent(self, service_account_key_file_json):
+        def set_initial_token(auth):
+            auth.initial_token = "test-initial-token"
+
+        valid = jwt.encode({"exp": 4102444800}, "secret", algorithm="HS256")
+        with (
+            patch.object(KeyAuth, "_KeyAuth__create_initial_token", new=set_initial_token),
+            patch("requests.post") as mock_post,
+            patch("stackit.core.auth_methods.key_auth.atexit") as mock_atexit,
+        ):
+            auth = self._make_auth(service_account_key_file_json, mock_post, valid)
+
+            auth.close()
+            auth.close()
+
+            assert not auth._refresh_thread.is_alive()
+            assert auth.executor._shutdown
+            mock_atexit.unregister.assert_called_with(auth._KeyAuth__shutdown)