diff --git a/components/runners/ambient-runner/ambient_runner/_grpc_client.py b/components/runners/ambient-runner/ambient_runner/_grpc_client.py index 2f7782af2..a5cacb98b 100644 --- a/components/runners/ambient-runner/ambient_runner/_grpc_client.py +++ b/components/runners/ambient-runner/ambient_runner/_grpc_client.py @@ -11,10 +11,11 @@ from pathlib import Path from typing import Optional +import grpc from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import padding -import grpc +from ambient_runner.platform.utils import set_bot_token logger = logging.getLogger(__name__) @@ -95,6 +96,7 @@ def _fetch_token_from_cp(cp_token_url: str, public_key_pem: str, session_id: str if not token: raise RuntimeError("CP /token response missing 'token' field") logger.info("[GRPC CLIENT] Fetched fresh API token from CP token endpoint") + set_bot_token(token) return token except urllib.error.HTTPError as e: resp_body = "" diff --git a/components/runners/ambient-runner/ambient_runner/platform/utils.py b/components/runners/ambient-runner/ambient_runner/platform/utils.py index 5f3393d22..6820c5623 100644 --- a/components/runners/ambient-runner/ambient_runner/platform/utils.py +++ b/components/runners/ambient-runner/ambient_runner/platform/utils.py @@ -23,14 +23,27 @@ # Kubelet automatically refreshes this file when the Secret is updated. _BOT_TOKEN_FILE = Path("/var/run/secrets/ambient/bot-token") +# In-process cache for the token fetched from the CP token endpoint. +# Set once at startup by _grpc_client.py after a successful CP token fetch. +_cp_fetched_token: str = "" + + +def set_bot_token(token: str) -> None: + """Store a token fetched from the CP token endpoint for use by get_bot_token().""" + global _cp_fetched_token + _cp_fetched_token = token.strip() + def get_bot_token() -> str: - """Return the current BOT_TOKEN, preferring the file mount over env var. + """Return the current BOT_TOKEN. - The operator mounts the runner-token Secret as a file so kubelet refreshes - it automatically when the token is rotated. Falls back to the BOT_TOKEN - env var for backward-compatibility with local / non-Kubernetes runs. + Priority: + 1. Token fetched from CP token endpoint (set via set_bot_token()). + 2. File mount at _BOT_TOKEN_FILE (kubelet-refreshed Secret). + 3. BOT_TOKEN env var (local / non-Kubernetes fallback). """ + if _cp_fetched_token: + return _cp_fetched_token try: if _BOT_TOKEN_FILE.exists(): return _BOT_TOKEN_FILE.read_text().strip() diff --git a/components/runners/ambient-runner/tests/test_grpc_client.py b/components/runners/ambient-runner/tests/test_grpc_client.py index baf379fd2..de6734307 100644 --- a/components/runners/ambient-runner/tests/test_grpc_client.py +++ b/components/runners/ambient-runner/tests/test_grpc_client.py @@ -200,6 +200,51 @@ def test_raises_on_missing_token_field(self): ) +class TestSetBotTokenIntegration: + def test_get_bot_token_returns_cp_fetched_token_after_successful_fetch(self): + import ambient_runner.platform.utils as utils + utils._cp_fetched_token = "" + + _, _, public_pem = generate_keypair() + mock_resp = MagicMock() + mock_resp.read.return_value = json.dumps({"token": "oidc-token-for-api-calls"}).encode() + mock_resp.__enter__ = MagicMock(return_value=mock_resp) + mock_resp.__exit__ = MagicMock(return_value=False) + + assert utils.get_bot_token() == "", "get_bot_token() must be empty before any CP fetch" + + with patch("urllib.request.urlopen", return_value=mock_resp): + _fetch_token_from_cp("http://cp.svc:8080/token", public_pem, "session-12345678") + + assert utils.get_bot_token() == "oidc-token-for-api-calls", ( + "get_bot_token() must return the CP-fetched token so backend API credential " + "calls are authenticated — regression for HTTP 401 on credential refresh" + ) + utils._cp_fetched_token = "" + + def test_fetch_from_cp_calls_set_bot_token(self): + from cryptography.hazmat.primitives.asymmetric import rsa as _rsa + private_key = _rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_pem = private_key.public_key().public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo, + ).decode() + + mock_resp = MagicMock() + mock_resp.read.return_value = json.dumps({"token": "oidc-api-token-abc"}).encode() + mock_resp.__enter__ = MagicMock(return_value=mock_resp) + mock_resp.__exit__ = MagicMock(return_value=False) + + import ambient_runner.platform.utils as utils + utils._cp_fetched_token = "" + + with patch("urllib.request.urlopen", return_value=mock_resp): + _fetch_token_from_cp("http://cp.svc:8080/token", public_pem, "session-12345678") + + assert utils.get_bot_token() == "oidc-api-token-abc" + utils._cp_fetched_token = "" + + class TestFromEnvIntegration: def test_uses_encrypted_session_id_when_cp_token_url_set(self): _, _, public_pem = generate_keypair()