Two-token mechanism for task execution to prevent token expiration while tasks wait in executor queues (#60108)

anishgirianish · web-flow · commit 2b6e8181e3ae · 2026-04-23T22:41:52.000+01:00
Tasks waiting in executor queues (Celery, Kubernetes) can have their JWT tokens expire before execution starts, causing auth failures on the Execution API. This is a real problem in production, when queues back up or workers are slow to pick up tasks, the original short-lived token expires and the worker gets a 403 when it finally tries to start the task. Fixes: #53713 Related: #59553 closes: #62129 ## Approach Two-token mechanism: a workload token (lifetime tracks [scheduler] task_queued_timeout) travels with the task through the queue, and a short-lived execution token is issued when the task actually starts running. The workload token carries a scope: "workload" claim and is restricted to the /run endpoint only, enforced via FastAPI SecurityScopes and a custom ExecutionAPIRoute. When /run succeeds, it returns an execution token via Refreshed-API-Token header. The SDK client picks it up and uses it for all subsequent API calls. The existing JWTReissueMiddleware handles refreshing execution tokens near expiry and skips workload tokens. Built on @ashb's SecurityScopes foundation. ### Security considerations Even if a workload token is intercepted, it can only call /run which already guards against running a task more than once (returns 409 if the task isn't in QUEUED/RESTARTING state). All other endpoints reject workload tokens , they require execution scope. The execution token issued by /run is short-lived and automatically refreshed, keeping the existing security posture for all API calls during task execution.
diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst
@@ -201,16 +201,25 @@ Token structure (Execution API)
 Token scopes (Execution API)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-The Execution API defines two token scopes:
+The Execution API defines two token scopes with different lifetimes:
 
 **workload**
-   A restricted scope accepted only on endpoints that explicitly opt in via
-   ``Security(require_auth, scopes=["token:workload"])``. Used for endpoints that
-   manage task state transitions.
+   A token embedded in the workload JSON payload when the Scheduler
+   dispatches a task. The longer lifetime
+   allows tasks to remain valid while waiting in executor queues before execution
+   begins. When a worker calls the ``/run`` endpoint with a ``workload`` token, the
+   server issues a fresh ``execution``-scoped token in the ``Refreshed-API-Token``
+   response header. Lifetime equals ``[scheduler] task_queued_timeout`` (default
+   600 seconds) — the same timeout the scheduler uses to reap queue-starved tasks —
+   so tuning ``task_queued_timeout`` also widens the window a task can wait in a
+   backed-up queue before its workload token expires.
 
 **execution**
-   Accepted by all Execution API endpoints. This is the standard scope for worker
-   communication and allows access
+   A short-lived token (default 10 minutes) accepted by all Execution API endpoints.
+   This is the standard scope for worker communication during task execution. Issued
+   by the server when the worker transitions to running via the ``/run`` endpoint.
+   The ``JWTReissueMiddleware`` refreshes ``execution`` tokens transparently,
+   so the worker maintains access for the duration of the task.
 
 Tokens without a ``scope`` claim default to ``"execution"`` for backwards compatibility.
 
@@ -219,14 +228,19 @@ Token delivery to workers
 
 The token flows through the execution stack as follows:
 
-1. **Scheduler** generates the token and embeds it in the workload JSON payload that it passes to
-   **Executor**.
+1. **Scheduler** generates a ``workload``-scoped token (lifetime equals
+   ``[scheduler] task_queued_timeout``, default 600 seconds) and embeds it in the workload
+   JSON payload that it passes to **Executor**.
 2. The workload JSON is passed to the worker process (via the executor-specific mechanism:
    Celery message, Kubernetes Pod spec, local subprocess arguments, etc.).
 3. The worker's ``execute_workload()`` function reads the workload JSON and extracts the token.
 4. The ``supervise()`` function receives the token and creates an ``httpx.Client`` instance
    with ``BearerAuth(token)`` for all Execution API HTTP requests.
-5. The token is included in the ``Authorization: Bearer <token>`` header of every request.
+5. The worker calls the ``/run`` endpoint with the ``workload``-scoped token to mark the task
+   as running. The server responds with a fresh ``execution``-scoped token in the
+   ``Refreshed-API-Token`` header.
+6. The client's ``_update_auth()`` hook detects the header and transparently updates
+   the ``BearerAuth`` instance to use the new ``execution`` token for all subsequent requests.
 
 Token validation (Execution API)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -251,7 +265,8 @@ Route-level enforcement is handled by ``require_auth``:
 Token refresh (Execution API)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-The ``JWTReissueMiddleware`` automatically refreshes valid tokens that are approaching expiry:
+The ``JWTReissueMiddleware`` automatically refreshes valid tokens that are approaching
+expiry. The token must be valid at the start of the request for refresh to occur:
 
 1. After each response, the middleware checks the token's remaining validity.
 2. If less than **20%** of the total validity remains (minimum 30 seconds), the server
@@ -260,16 +275,20 @@ The ``JWTReissueMiddleware`` automatically refreshes valid tokens that are appro
 4. The client's ``_update_auth()`` hook detects this header and transparently updates
    the ``BearerAuth`` instance for subsequent requests.
 
-This mechanism ensures long-running tasks do not lose API access due to token expiry,
-without requiring the worker to re-authenticate.
+The middleware only refreshes ``execution``-scoped tokens. ``workload``-scoped tokens are
+sized to span the queued-timeout window and are explicitly skipped by the middleware —
+they are designed to survive executor queue wait times without needing refresh. This
+ensures long-running tasks do not lose API access without requiring the worker to
+re-authenticate.
 
 No token revocation (Execution API)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Execution API tokens are not subject to revocation. They are short-lived (default 10 minutes)
-and automatically refreshed by the ``JWTReissueMiddleware``, so revocation is not part of the
-Execution API security model. Once an Execution API token is issued to a worker, it remains
-valid until it expires.
+Execution API tokens are not subject to revocation. ``execution``-scoped tokens are short-lived
+(default 10 minutes) and automatically refreshed by the ``JWTReissueMiddleware``.
+``workload``-scoped tokens (tracking ``[scheduler] task_queued_timeout``) are not refreshed —
+they expire naturally after their validity period. Revocation is not part of the Execution API
+security model.
 
 
 
@@ -284,11 +303,12 @@ Default timings (Execution API)
      - Default
    * - ``[execution_api] jwt_expiration_time``
      - 600 seconds (10 minutes)
+   * - Workload token lifetime (derived)
+     - ``[scheduler] task_queued_timeout`` (default 600 seconds)
    * - ``[execution_api] jwt_audience``
      - ``urn:airflow.apache.org:task``
    * - Token refresh threshold
-     - 20% of validity remaining (minimum 30 seconds, i.e., at ~120 seconds before expiry
-       with the default 600-second token lifetime)
+     - 20% of validity remaining (minimum 30 seconds)
 
 
 Dag File Processor and Triggerer
@@ -386,7 +406,10 @@ All JWT-related configuration parameters:
      - JWKS endpoint URL or local file path for token validation. Mutually exclusive with ``jwt_secret``.
    * - ``[execution_api] jwt_expiration_time``
      - 600 (10 min)
-     - Execution API token lifetime in seconds.
+     - Execution API ``execution``-scoped token lifetime in seconds.
+   * - ``[scheduler] task_queued_timeout``
+     - 600.0 (10 min)
+     - Queue-starvation timeout. Also sets the ``workload``-scoped token lifetime to the same value.
    * - ``[execution_api] jwt_audience``
      - ``urn:airflow.apache.org:task``
      - Audience claim for Execution API tokens.
diff --git a/airflow-core/src/airflow/api_fastapi/auth/tokens.py b/airflow-core/src/airflow/api_fastapi/auth/tokens.py
@@ -447,15 +447,21 @@ def signing_arg(self) -> AllowedPrivateKeys | str:
             assert self._secret_key
         return self._secret_key
 
-    def generate(self, extras: dict[str, Any] | None = None, headers: dict[str, Any] | None = None) -> str:
+    def generate(
+        self,
+        extras: dict[str, Any] | None = None,
+        headers: dict[str, Any] | None = None,
+        valid_for: float | None = None,
+    ) -> str:
         """Generate a signed JWT for the subject."""
         now = int(datetime.now(tz=timezone.utc).timestamp())
+        effective_valid_for = valid_for if valid_for is not None else self.valid_for
         claims = {
             "jti": uuid.uuid4().hex,
             "iss": self.issuer,
             "aud": self.audience,
             "nbf": now,
-            "exp": int(now + self.valid_for),
+            "exp": int(now + effective_valid_for),
             "iat": now,
         }
 
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py
@@ -129,8 +129,6 @@ async def dispatch(self, request: Request, call_next):
 
 class JWTReissueMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):
-        from airflow.configuration import conf
-
         response: Response = await call_next(request)
 
         refreshed_token: str | None = None
@@ -142,9 +140,15 @@ async def dispatch(self, request: Request, call_next):
                     validator: JWTValidator = await services.aget(JWTValidator)
                     claims = await validator.avalidated_claims(token, {})
 
+                    # Workload tokens are long-lived and meant to survive queue
+                    # wait times so avoid refreshing them. If avalidated_claims
+                    # raises for a workload token, the outer except handles it.
+                    if claims.get("scope") == "workload":
+                        return response
+
                     now = int(time.time())
-                    validity = conf.getint("execution_api", "jwt_expiration_time")
-                    refresh_when_less_than = max(int(validity * 0.20), 30)
+                    token_lifetime = int(claims.get("exp", 0)) - int(claims.get("iat", 0))
+                    refresh_when_less_than = max(int(token_lifetime * 0.20), 30)
                     valid_left = int(claims.get("exp", 0)) - now
                     if valid_left <= refresh_when_less_than:
                         generator: JWTGenerator = await services.aget(JWTGenerator)
@@ -312,7 +316,6 @@ class InProcessExecutionAPI:
     def app(self):
         if not self._app:
             from airflow.api_fastapi.common.dagbag import create_dag_bag
-            from airflow.api_fastapi.execution_api.app import create_task_execution_api_app
             from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
             from airflow.api_fastapi.execution_api.routes.connections import has_connection_access
             from airflow.api_fastapi.execution_api.routes.variables import has_variable_access
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -28,7 +28,7 @@
 import attrs
 import structlog
 from cadwyn import VersionedAPIRouter
-from fastapi import Body, HTTPException, Query, Security, status
+from fastapi import Body, HTTPException, Query, Response, Security, status
 from opentelemetry import trace
 from opentelemetry.trace import StatusCode
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
@@ -42,6 +42,7 @@
 
 from airflow._shared.observability.traces import override_ids
 from airflow._shared.timezones import timezone
+from airflow.api_fastapi.auth.tokens import JWTGenerator
 from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag
 from airflow.api_fastapi.common.db.common import SessionDep
 from airflow.api_fastapi.common.types import UtcDateTime
@@ -63,7 +64,9 @@
     TISuccessStatePayload,
     TITerminalStatePayload,
 )
-from airflow.api_fastapi.execution_api.security import ExecutionAPIRoute, require_auth
+from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.deps import DepContainer
+from airflow.api_fastapi.execution_api.security import CurrentTIToken, ExecutionAPIRoute, require_auth
 from airflow.exceptions import TaskNotFound
 from airflow.models.asset import AssetActive
 from airflow.models.dag import DagModel
@@ -98,6 +101,7 @@
 @ti_id_router.patch(
     "/{task_instance_id}/run",
     status_code=status.HTTP_200_OK,
+    dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])],
     responses={
         status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
         status.HTTP_409_CONFLICT: {"description": "The TI is already in the requested state"},
@@ -108,8 +112,11 @@
 def ti_run(
     task_instance_id: UUID,
     ti_run_payload: Annotated[TIEnterRunningPayload, Body()],
+    response: Response,
     session: SessionDep,
     dag_bag: DagBagDep,
+    services=DepContainer,
+    token: TIToken = CurrentTIToken,
 ) -> TIRunContext:
     """
     Run a TaskInstance.
@@ -293,13 +300,20 @@ def ti_run(
             context.next_method = ti.next_method
             context.next_kwargs = ti.next_kwargs
             context.start_date = ti.start_date
-        return context
     except SQLAlchemyError:
         log.exception("Error marking Task Instance state as running")
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Database error occurred"
         )
 
+    # JWTReissueMiddleware also writes Refreshed-API-Token but skips workload tokens, so we set it here for the workload→execution swap.
+    if token.claims.scope == "workload":
+        generator: JWTGenerator = services.get(JWTGenerator)
+        execution_token = generator.generate(extras={"sub": str(task_instance_id), "scope": "execution"})
+        response.headers["Refreshed-API-Token"] = execution_token
+
+    return context
+
 
 @ti_id_router.patch(
     "/{task_instance_id}/state",
diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml
@@ -2586,6 +2586,10 @@ scheduler:
     task_queued_timeout:
       description: |
         Amount of time a task can be in the queued state before being retried or set to failed.
+
+        This value also sets the lifetime of the workload JWT token that is sent with the task
+        to an executor queue, so a task waiting in the queue can still authenticate to the
+        Execution API until its queue-starvation deadline.
       version_added: 2.6.0
       type: float
       example: ~
diff --git a/airflow-core/src/airflow/executors/workloads/base.py b/airflow-core/src/airflow/executors/workloads/base.py
@@ -25,6 +25,8 @@
 
 from pydantic import BaseModel, ConfigDict, Field
 
+from airflow.configuration import conf
+
 if TYPE_CHECKING:
     from airflow.api_fastapi.auth.tokens import JWTGenerator
     from airflow.executors.workloads.types import WorkloadState
@@ -76,7 +78,13 @@ class BaseWorkloadSchema(BaseModel):
 
     @staticmethod
     def generate_token(sub_id: str, generator: JWTGenerator | None = None) -> str:
-        return generator.generate({"sub": sub_id}) if generator else ""
+        if not generator:
+            return ""
+        valid_for = conf.getfloat("scheduler", "task_queued_timeout")
+        return generator.generate(
+            extras={"sub": sub_id, "scope": "workload"},
+            valid_for=valid_for,
+        )
 
 
 class BaseDagBundleWorkload(BaseWorkloadSchema, ABC):
diff --git a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py
@@ -160,6 +160,34 @@ def test_secret_key_with_configured_kid():
         assert header["kid"] == "my-custom-kid"
 
 
+def test_generate_with_custom_valid_for():
+    """generate() accepts a valid_for override."""
+    generator = JWTGenerator(secret_key="test-secret", audience="test", valid_for=60)
+    token = generator.generate(extras={"sub": "user"}, valid_for=3600)
+    claims = jwt.decode(token, "test-secret", algorithms=["HS512"], audience="test")
+    assert claims["exp"] - claims["iat"] == 3600
+
+
+def test_generate_workload_scope_via_extras():
+    """generate() with scope='workload' in extras produces a workload-scoped token."""
+    generator = JWTGenerator(secret_key="test-secret", audience="test", valid_for=60)
+
+    token = generator.generate(extras={"sub": "ti-123", "scope": "workload"}, valid_for=86400)
+    claims = jwt.decode(token, "test-secret", algorithms=["HS512"], audience="test")
+    assert claims["sub"] == "ti-123"
+    assert claims["scope"] == "workload"
+    assert claims["exp"] - claims["iat"] == 86400
+
+
+def test_regular_token_has_no_scope():
+    """Regular tokens without scope in extras have no scope claim."""
+    generator = JWTGenerator(secret_key="test-secret", audience="test", valid_for=60)
+
+    regular = generator.generate(extras={"sub": "user"})
+    regular_claims = jwt.decode(regular, "test-secret", algorithms=["HS512"], audience="test")
+    assert "scope" not in regular_claims
+
+
 @pytest.fixture
 def jwt_generator(ed25519_private_key: Ed25519PrivateKey):
     key = ed25519_private_key
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py b/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py
@@ -22,8 +22,16 @@
 from starlette.routing import Mount
 
 from airflow.api_fastapi.app import cached_app
+from airflow.api_fastapi.execution_api.app import lifespan
 from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
-from airflow.api_fastapi.execution_api.security import _jwt_bearer
+from airflow.api_fastapi.execution_api.security import require_auth
+
+
+@pytest.fixture(autouse=True)
+def _restore_lifespan_registry():
+    snapshot = dict(lifespan.registry._services)
+    yield
+    lifespan.registry._services = snapshot
 
 
 def _get_execution_api_app(root_app: FastAPI) -> FastAPI:
@@ -45,16 +53,15 @@ def client(request: pytest.FixtureRequest):
     app = cached_app(apps="execution")
     exec_app = _get_execution_api_app(app)
 
-    async def mock_jwt_bearer(request: Request):
+    async def mock_require_auth(request: Request) -> TIToken:
         from uuid import UUID
 
         ti_id = UUID(request.path_params.get("task_instance_id", "00000000-0000-0000-0000-000000000000"))
-        claims = TIClaims(scope="execution")
-        return TIToken(id=ti_id, claims=claims)
+        return TIToken(id=ti_id, claims=TIClaims(scope="execution"))
 
-    exec_app.dependency_overrides[_jwt_bearer] = mock_jwt_bearer
+    exec_app.dependency_overrides[require_auth] = mock_require_auth
 
     with TestClient(app, headers={"Authorization": "Bearer fake"}) as client:
         yield client
 
-    exec_app.dependency_overrides.pop(_jwt_bearer, None)
+    exec_app.dependency_overrides.pop(require_auth, None)
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_router.py
@@ -25,8 +25,6 @@
 from airflow.api_fastapi.auth.tokens import JWTValidator
 from airflow.api_fastapi.execution_api.app import lifespan
 
-from tests_common.test_utils.config import conf_vars
-
 
 @pytest.fixture
 def exec_app(client):
@@ -53,6 +51,7 @@ def test_expiring_token_is_reissued(
     auth = AsyncMock(spec=JWTValidator)
     auth.avalidated_claims.return_value = {
         "sub": "edb09971-4e0e-4221-ad3f-800852d38085",
+        "iat": moment,
         "exp": moment + validity,
     }
 
@@ -62,8 +61,7 @@ def test_expiring_token_is_reissued(
     lifespan.registry.register_value(JWTValidator, auth)
     # In order to test this we need any endpoint to hit. The easiest one to use is variable get
 
-    with conf_vars({("execution_api", "jwt_expiration_time"): str(validity)}):
-        response = client.get("/execution/variables/key1", headers={"Authorization": "Bearer dummy"})
+    response = client.get("/execution/variables/key1", headers={"Authorization": "Bearer dummy"})
 
     if expect_refreshed_token:
         assert "Refreshed-API-Token" in response.headers
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py
diff --git a/devel-common/src/tests_common/test_utils/mock_executor.py b/devel-common/src/tests_common/test_utils/mock_executor.py
