Fix Dag processor identity exchange edge cases

Processors must stop trusting keys withdrawn by their issuer and recover from temporary identity-provider failures without exhausting retries during its cooldown. Authentication must also remain the boundary for incoming trace context.

Author: ephraimbuddy

airflow-core/src/airflow/api_fastapi/auth/tokens.py

@@ -163,18 +163,22 @@ async def fetch_jwks(self, *, force: bool = False) -> None:
                 return
 
             try:
-                # PyJWKSet.from_dict raises AttributeError for a JSON array or scalar.
+                # Older PyJWT versions assume the body and every key are objects.
                 if not isinstance(data, dict):
                     raise TypeError(f"JWKS must be a JSON object, not {type(data).__name__}")
-                keyset = jwt.PyJWKSet.from_dict(data)
+                keys = data.get("keys")
+                if not isinstance(keys, list) or any(not isinstance(key, dict) for key in keys):
+                    raise TypeError("JWKS keys must be an array of JSON objects")
+                # PyJWT rejects empty sets, but an issuer can remove all keys with an empty array.
+                keyset = jwt.PyJWKSet.from_dict(data) if keys else None
             except (jwt.PyJWKError, jwt.PyJWKSetError, KeyError, TypeError, ValueError):
                 if not self.raise_on_fetch_error:
                     raise
                 log.warning("Unable to load JWKS response", url=self.url, exc_info=True)
                 return
             self._jwks = keyset
             self._fetch_failed = False
-            log.debug("Fetched JWKS", url=self.url, keys=len(self._jwks.keys))
+            log.debug("Fetched JWKS", url=self.url, keys=len(keyset.keys) if keyset is not None else 0)
 
     async def _fetch_remote_jwks(self) -> dict[str, Any] | None:
         try:
@@ -251,10 +255,11 @@ async def get_key(self, kid: str) -> jwt.PyJWK:
             except KeyError:
                 if self.refresh_on_unknown_kid:
                     await self.fetch_jwks(force=True)
-                    try:
-                        return self._jwks[kid]
-                    except KeyError:
-                        pass
+                    if self._jwks is not None:
+                        try:
+                            return self._jwks[kid]
+                        except KeyError:
+                            pass
 
         if self.raise_on_fetch_error and self._fetch_failed:
             raise JWKSUnavailableError("Unable to load the required signing key")

airflow-core/src/airflow/api_fastapi/execution_api/app.py

@@ -292,9 +292,13 @@ def _inject_trace_context_dep(routes, mode: str) -> None:
             case "unsafe-always":
                 route.dependencies.insert(0, dep)
             case "only-authenticated":
+                from airflow.api_fastapi.execution_api.routes.session_tokens import require_workload_identity
                 from airflow.api_fastapi.execution_api.security import require_auth
 
-                if any(getattr(d, "dependency", None) is require_auth for d in route.dependencies):
+                if any(
+                    getattr(d, "dependency", None) in (require_auth, require_workload_identity)
+                    for d in route.dependencies
+                ):
                     route.dependencies.append(dep)
 
 

airflow-core/src/airflow/api_fastapi/execution_api/routes/session_tokens.py

@@ -76,6 +76,7 @@ async def require_workload_identity(
 
 @router.post(
     "/dag-processor/session-token",
+    dependencies=[Depends(require_workload_identity)],
     responses=create_openapi_http_exception_doc(
         [
             (status.HTTP_401_UNAUTHORIZED, "The workload identity is missing, invalid, or expired"),

airflow-core/src/airflow/api_fastapi/execution_api/workload_identity.py

@@ -26,7 +26,7 @@
 from typing import Annotated, Literal
 from uuid import NAMESPACE_URL, UUID, uuid5
 
-import httpx
+import httpx2
 import jwt
 import structlog
 from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator
@@ -84,7 +84,7 @@ class _IdentityClaims(BaseModel):
 class DagProcessorIdentityVerifier:
     """Validate external JWTs independently of Airflow's runtime token trust."""
 
-    def __init__(self, config: WorkloadIdentityConfig, *, client: httpx.AsyncClient):
+    def __init__(self, config: WorkloadIdentityConfig, *, client: httpx2.AsyncClient):
         if config.audience in conf.getlist("execution_api", "jwt_audience"):
             raise ValueError("Workload identity must use a dedicated audience")
         configured = get_configured_bundle_names()
@@ -163,5 +163,5 @@ async def create_dag_processor_identity_verifier() -> AsyncIterator[DagProcessor
         yield None
         return
     config = WorkloadIdentityConfig.model_validate_json(raw)
-    async with httpx.AsyncClient(timeout=5.0, follow_redirects=False) as client:
+    async with httpx2.AsyncClient(timeout=5.0, follow_redirects=False) as client:
         yield DagProcessorIdentityVerifier(config, client=client)

airflow-core/src/airflow/dag_processing/credentials.py

@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+import math
 from collections.abc import Callable
 from pathlib import Path
 from time import monotonic
@@ -55,7 +56,9 @@ def _wait_for_retry_after(retry_state: RetryCallState) -> float:
         retry_after = float(error.response.headers["Retry-After"])
     except (KeyError, ValueError):  # Missing, or an HTTP-date rather than seconds.
         return wait
-    return max(wait, min(retry_after, API_RETRY_WAIT_MAX))
+    if not math.isfinite(retry_after):
+        return wait
+    return max(wait, retry_after)
 
 
 class SessionTokenProvider(Protocol):

airflow-core/tests/unit/api_fastapi/auth/test_tokens.py

@@ -70,9 +70,9 @@ async def test_refresh_unknown_key_before_regular_interval(
 
         def handle(request):
             requests.append(request)
-            return httpx.Response(200, json={"keys": keys})
+            return httpx2.Response(200, json={"keys": keys})
 
-        async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client:
+        async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client:
             jwks = JWKS(
                 url="https://example.com/keys", client=client, refresh_on_unknown_kid=refresh_on_unknown_kid
             )
@@ -98,9 +98,9 @@ async def test_concurrent_unknown_keys_share_fetch_and_cooldown(
         async def handle(request):
             requests.append(request)
             await asyncio.sleep(0)
-            return httpx.Response(status, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]})
+            return httpx2.Response(status, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]})
 
-        async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client:
+        async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client:
             jwks = JWKS(
                 url="https://example.com/keys",
                 client=client,
@@ -135,9 +135,9 @@ async def test_failed_unknown_key_refresh_keeps_cached_keys_and_rate_limits(
 
         def handle(request):
             requests.append(request)
-            return httpx.Response(status, json={"keys": keys})
+            return httpx2.Response(status, json={"keys": keys})
 
-        async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client:
+        async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client:
             jwks = JWKS(
                 url="https://example.com/keys",
                 client=client,
@@ -172,9 +172,9 @@ async def handle(request):
             if len(requests) > 1:
                 fetching.set()
                 await release.wait()
-            return httpx.Response(200, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]})
+            return httpx2.Response(200, json={"keys": [key_to_jwk_dict(ed25519_private_key, "known")]})
 
-        async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client:
+        async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client:
             jwks = JWKS(url="https://example.com/keys", client=client, refresh_on_unknown_kid=True)
             known = await jwks.get_key("known")
             clock.return_value = 111
@@ -195,14 +195,15 @@ async def handle(request):
                 await refresh
             assert len(requests) == 2
 
+    @pytest.mark.parametrize("body", [[{"kty": "OKP"}], {"keys": [None]}, {"keys": ["invalid"]}])
     @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100)
-    async def test_jwks_body_that_is_not_an_object_is_a_failed_fetch(self, clock, ed25519_private_key):
-        bodies = [[{"kty": "OKP"}]]
+    async def test_malformed_jwks_is_a_failed_fetch(self, clock, ed25519_private_key, body):
+        bodies = [body]
 
         def handle(request):
-            return httpx.Response(200, json=bodies[-1])
+            return httpx2.Response(200, json=bodies[-1])
 
-        async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as client:
+        async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client:
             jwks = JWKS(
                 url="https://example.com/keys",
                 client=client,
@@ -216,9 +217,37 @@ def handle(request):
             known = await jwks.get_key("known")
             # Once the hourly refresh is due, a bad body must not stop the cached key from validating.
             clock.return_value = 111 + jwks.refresh_interval_secs + 1
-            bodies.append([{"kty": "OKP"}])
+            bodies.append(body)
             assert await jwks.get_key("known") is known
 
+    @pytest.mark.parametrize("refresh", ["initial", "scheduled", "unknown-key"])
+    @patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100)
+    async def test_empty_jwks_removes_cached_keys(self, clock, ed25519_private_key, refresh):
+        keys = [] if refresh == "initial" else [key_to_jwk_dict(ed25519_private_key, "known")]
+        requests = []
+
+        def handle(request):
+            requests.append(request)
+            return httpx2.Response(200, json={"keys": keys})
+
+        async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client:
+            jwks = JWKS(
+                url="https://example.com/keys",
+                client=client,
+                refresh_on_unknown_kid=True,
+                raise_on_fetch_error=True,
+            )
+            if refresh != "initial":
+                await jwks.get_key("known")
+                keys.clear()
+                clock.return_value += jwks.refresh_interval_secs + 1 if refresh == "scheduled" else 11
+            if refresh == "unknown-key":
+                with pytest.raises(KeyError):
+                    await jwks.get_key("unknown")
+            with pytest.raises(KeyError):
+                await jwks.get_key("known")
+            assert len(requests) == (1 if refresh == "initial" else 2)
+
     @pytest.mark.parametrize("private_key", ["rsa_private_key", "ed25519_private_key"], indirect=True)
     async def test_fetch_jwks_success(self, private_key):
         jwk_content = json.dumps({"keys": [key_to_jwk_dict(private_key, "kid")]})

airflow-core/tests/unit/api_fastapi/execution_api/test_app.py

@@ -21,10 +21,12 @@
 import json
 import threading
 from contextlib import asynccontextmanager
+from datetime import UTC, datetime
 from unittest import mock
 from uuid import UUID
 
 import httpx2
+import jwt
 import pytest
 from fastapi import FastAPI, Request, status
 from fastapi.params import Security as SecurityParam
@@ -46,6 +48,10 @@
 from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
 from airflow.api_fastapi.execution_api.security import require_auth
 from airflow.api_fastapi.execution_api.versions import bundle
+from airflow.api_fastapi.execution_api.workload_identity import (
+    DagProcessorIdentity,
+    DagProcessorIdentityVerifier,
+)
 from airflow.models.xcom import XComModel
 from airflow.utils.session import create_session_async
 
@@ -477,6 +483,47 @@ async def mock_require_auth(request: Request) -> TIToken:
         if expect_status is not None:
             assert response.status_code == expect_status
 
+    @pytest.mark.parametrize("mode", ["only-authenticated", "unsafe-always", "never"])
+    @pytest.mark.parametrize("valid_identity", [True, False])
+    @mock.patch.object(otel_propagate, "extract", autospec=True)
+    def test_workload_identity_trace_context_extraction(self, extract, mode, valid_identity, time_machine):
+        now = datetime(2026, 10, 5, 12, 0, tzinfo=UTC)
+        time_machine.move_to(now, tick=False)
+        events = []
+        verifier = mock.create_autospec(DagProcessorIdentityVerifier, instance=True)
+
+        async def authenticate(token):
+            events.append("authenticate")
+            if not valid_identity:
+                raise jwt.InvalidTokenError("Invalid identity")
+            return DagProcessorIdentity(UUID(int=1), frozenset({"dags-folder"}), now.timestamp() + 600)
+
+        def extract_context(headers):
+            events.append("extract")
+            return otel_context.get_current()
+
+        verifier.authenticate.side_effect = authenticate
+        extract.side_effect = extract_context
+        app = self._build_app(mode)
+        with TestClient(app) as client:
+            app.state.svcs_registry.register_value(DagProcessorIdentityVerifier, verifier)
+            response = client.post(
+                "/dag-processor/session-token",
+                headers={
+                    "Authorization": "Bearer workload-identity",
+                    "traceparent": "00-12345678901234567890123456789012-1234567890123456-01",
+                },
+            )
+
+        assert response.status_code == (200 if valid_identity else 401)
+        expect_extract = mode == "unsafe-always" or (mode == "only-authenticated" and valid_identity)
+        assert extract.call_count == int(expect_extract)
+        assert events[0] == ("extract" if mode == "unsafe-always" else "authenticate")
+        if expect_extract:
+            assert extract.call_args.args[0]["traceparent"] == (
+                "00-12345678901234567890123456789012-1234567890123456-01"
+            )
+
     def test_trace_context_dep_cleans_up_on_route_exception(self):
         """Verify extract and cleanup run correctly when a route handler raises."""
         app = self._build_app("unsafe-always")

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_session_tokens.py

@@ -20,7 +20,6 @@
 from datetime import UTC, datetime
 from unittest import mock
 
-import httpx
 import httpx2
 import jwt
 import pytest
@@ -125,8 +124,8 @@ def verifier(client, exec_app, signing_key, identity_algorithm, monkeypatch):
         }
     else:
         public_key = key_to_jwk_dict(signing_key, kid="key-1")
-    jwks_client = httpx.AsyncClient(
-        transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"keys": [public_key]}))
+    jwks_client = httpx2.AsyncClient(
+        transport=httpx2.MockTransport(lambda request: httpx2.Response(200, json={"keys": [public_key]}))
     )
     verifier = svcs.Container(exec_app.state.svcs_registry).get(DagProcessorIdentityVerifier)
     monkeypatch.setattr(verifier.validator.jwks, "client", jwks_client)
@@ -305,7 +304,7 @@ def test_rotated_signing_key_refreshes_after_cooldown(
 
     async def get_keys(url):
         requests.append(url)
-        return httpx.Response(200, json={"keys": keys}, request=httpx.Request("GET", url))
+        return httpx2.Response(200, json={"keys": keys}, request=httpx2.Request("GET", url))
 
     monkeypatch.setattr(verifier.validator.jwks.client, "get", get_keys)
     original = decode(exchange(client, identity_token()))
@@ -371,15 +370,15 @@ async def expire(token):
 def test_jwks_outage_is_retryable(client, verifier, identity_token, monkeypatch, failure):
     async def unavailable(url):
         if failure == "connection":
-            raise httpx.ConnectError("offline")
-        request = httpx.Request("GET", url)
+            raise httpx2.ConnectError("offline")
+        request = httpx2.Request("GET", url)
         if failure == "503":
-            return httpx.Response(503, request=request)
+            return httpx2.Response(503, request=request)
         if failure == "invalid-json":
-            return httpx.Response(200, content=b"unavailable", request=request)
+            return httpx2.Response(200, content=b"unavailable", request=request)
         if failure == "json-array":
-            return httpx.Response(200, json=[{"kty": "OKP"}], request=request)
-        return httpx.Response(200, json={"keys": []}, request=request)
+            return httpx2.Response(200, json=[{"kty": "OKP"}], request=request)
+        return httpx2.Response(200, json={"keys": [None]}, request=request)
 
     monkeypatch.setattr(verifier.jwks.client, "get", unavailable)
     response = exchange(client, identity_token())
@@ -388,12 +387,32 @@ async def unavailable(url):
     assert response.headers["Retry-After"] == "10"
 
 
+@pytest.mark.parametrize("cached", [False, True])
+@mock.patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100)
+def test_empty_jwks_rejects_identities_with_or_without_cached_keys(
+    clock, client, verifier, identity_token, monkeypatch, cached
+):
+    if cached:
+        assert exchange(client, identity_token()).status_code == 200
+        clock.return_value = 3701
+
+    async def get_keys(url):
+        return httpx2.Response(200, json={"keys": []}, request=httpx2.Request("GET", url))
+
+    monkeypatch.setattr(verifier.jwks.client, "get", get_keys)
+    response = exchange(client, identity_token())
+    assert response.status_code == 401
+    assert response.json() == {"detail": "Invalid workload identity"}
+    assert response.headers["WWW-Authenticate"] == "Bearer"
+    assert "Retry-After" not in response.headers
+
+
 @mock.patch("airflow.api_fastapi.auth.tokens.time.monotonic", autospec=True, return_value=100)
 def test_jwks_outage_preserves_cached_key_validation(clock, client, verifier, identity_token, monkeypatch):
     assert exchange(client, identity_token()).status_code == 200
 
     async def unavailable(url):
-        raise httpx.ConnectError("offline")
+        raise httpx2.ConnectError("offline")
 
     monkeypatch.setattr(verifier.jwks.client, "get", unavailable)
     clock.return_value = 3701
@@ -408,10 +427,10 @@ def test_registration_recovers_from_jwks_outage(clock, make_processor, verifier,
 
     async def get_keys(url):
         attempts.append(url)
-        return httpx.Response(
+        return httpx2.Response(
             503 if len(attempts) == 1 else 200,
             json={"keys": [key_to_jwk_dict(signing_key, kid="key-1")]},
-            request=httpx.Request("GET", url),
+            request=httpx2.Request("GET", url),
         )
 
     waits = []