further enhanced the implementation

Author: anishgirianish

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

@@ -46,7 +46,7 @@
     "JWKS",
     "JWTGenerator",
     "JWTValidator",
-    "TOKEN_SCOPE_QUEUE",
+    "TOKEN_SCOPE_WORKLOAD",
     "generate_private_key",
     "get_sig_validation_args",
     "get_signing_args",
@@ -55,7 +55,7 @@
     "key_to_jwk_dict",
 ]
 
-TOKEN_SCOPE_QUEUE = "queue"
+TOKEN_SCOPE_WORKLOAD = "ExecuteTaskWorkload"
 
 
 class InvalidClaimError(ValueError):
@@ -437,15 +437,28 @@ 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:
-        """Generate a signed JWT for the subject."""
+    def generate(
+        self,
+        extras: dict[str, Any] | None = None,
+        headers: dict[str, Any] | None = None,
+        expiry: int | None = None,
+    ) -> str:
+        """
+        Generate a signed JWT.
+
+        Args:
+            extras: Additional claims to include in the token. These are merged with default claims.
+            headers: Additional headers to include in the JWT.
+            expiry: Optional custom expiry time in seconds. If not provided, uses self.valid_for.
+        """
         now = int(datetime.now(tz=timezone.utc).timestamp())
+        valid_for = expiry if expiry 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 + valid_for),
             "iat": now,
         }
 
@@ -461,38 +474,20 @@ def generate(self, extras: dict[str, Any] | None = None, headers: dict[str, Any]
             headers["kid"] = self.kid
         return jwt.encode(claims, self.signing_arg, algorithm=self.algorithm, headers=headers)
 
-    def generate_queue_token(self, sub: str) -> str:
+    def generate_workload_token(self, sub: str) -> str:
         """
-        Generate a long-lived queue token for task workloads.
+        Generate a long-lived workload token for task execution.
 
-        Queue tokens have a special 'scope' claim that restricts them to the /run endpoint only.
-        They are valid for longer (default 24h) to survive queue wait times.
+        Workload tokens have a special 'scope' claim that restricts them to the /run endpoint only.
+        They are valid for longer (default 24h) to survive executor queue wait times.
         """
         from airflow.configuration import conf
 
-        queue_expiry = conf.getint("execution_api", "jwt_queue_token_expiration_time", fallback=86400)
-        now = int(datetime.now(tz=timezone.utc).timestamp())
-
-        claims = {
-            "jti": uuid.uuid4().hex,
-            "iss": self.issuer,
-            "aud": self.audience,
-            "nbf": now,
-            "exp": now + queue_expiry,
-            "iat": now,
-            "sub": sub,
-            "scope": TOKEN_SCOPE_QUEUE,
-        }
-
-        if claims["iss"] is None:
-            del claims["iss"]
-        if claims["aud"] is None:
-            del claims["aud"]
-
-        headers = {"alg": self.algorithm}
-        if self._private_key:
-            headers["kid"] = self.kid
-        return jwt.encode(claims, self.signing_arg, algorithm=self.algorithm, headers=headers)
+        workload_expiry = conf.getint("execution_api", "jwt_workload_token_expiration_time", fallback=86400)
+        return self.generate(
+            extras={"sub": sub, "scope": TOKEN_SCOPE_WORKLOAD},
+            expiry=workload_expiry,
+        )
 
 
 def generate_private_key(key_type: str = "RSA", key_size: int = 2048):

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

@@ -296,24 +296,16 @@ class InProcessExecutionAPI:
     @cached_property
     def app(self):
         if not self._app:
-            from unittest.mock import AsyncMock, MagicMock
-
-            from airflow.api_fastapi.auth.tokens import JWTValidator
             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.deps import (
-                DepContainer,
                 JWTBearerDep,
-                JWTBearerQueueDep,
                 JWTBearerTIPathDep,
             )
             from airflow.api_fastapi.execution_api.routes.connections import has_connection_access
+            from airflow.api_fastapi.execution_api.routes.task_instances import JWTBearerWorkloadDep
             from airflow.api_fastapi.execution_api.routes.variables import has_variable_access
             from airflow.api_fastapi.execution_api.routes.xcoms import has_xcom_access
-            from airflow.configuration import conf
-
-            # Set a dummy JWT secret so the lifespan can create JWT services without failing.
-            if not conf.get("api_auth", "jwt_secret", fallback=None):
-                conf.set("api_auth", "jwt_secret", "in-process-test-secret-key")
 
             self._app = create_task_execution_api_app()
 
@@ -324,33 +316,11 @@ async def always_allow(): ...
 
             self._app.dependency_overrides[JWTBearerDep.dependency] = always_allow
             self._app.dependency_overrides[JWTBearerTIPathDep.dependency] = always_allow
-            self._app.dependency_overrides[JWTBearerQueueDep.dependency] = always_allow
+            self._app.dependency_overrides[JWTBearerWorkloadDep.dependency] = always_allow
             self._app.dependency_overrides[has_connection_access] = always_allow
             self._app.dependency_overrides[has_variable_access] = always_allow
             self._app.dependency_overrides[has_xcom_access] = always_allow
 
-            # Create a mock container that provides mock JWT services
-            mock_jwt_generator = MagicMock(spec=JWTGenerator)
-            mock_jwt_generator.generate.return_value = "mock-execution-token"
-
-            mock_jwt_validator = AsyncMock(spec=JWTValidator)
-            mock_jwt_validator.avalidated_claims.return_value = {"sub": "test", "exp": 9999999999}
-
-            class MockContainer:
-                """A mock svcs container that returns mock services."""
-
-                async def aget(self, svc_type):
-                    if svc_type is JWTGenerator:
-                        return mock_jwt_generator
-                    if svc_type is JWTValidator:
-                        return mock_jwt_validator
-                    raise ValueError(f"Unknown service type: {svc_type}")
-
-            async def mock_container_dep():
-                return MockContainer()
-
-            self._app.dependency_overrides[DepContainer.dependency] = mock_container_dep
-
         return self._app
 
     @cached_property

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

@@ -26,7 +26,7 @@
 from fastapi.security import HTTPBearer
 from sqlalchemy import select
 
-from airflow.api_fastapi.auth.tokens import TOKEN_SCOPE_QUEUE, JWTValidator
+from airflow.api_fastapi.auth.tokens import TOKEN_SCOPE_WORKLOAD, JWTValidator
 from airflow.api_fastapi.common.db.common import AsyncSessionDep
 from airflow.api_fastapi.execution_api.datamodels.token import TIToken
 from airflow.configuration import conf
@@ -46,15 +46,12 @@ async def _container(request: Request):
 DepContainer: svcs.Container = Depends(_container)
 
 
-class JWTBearer(HTTPBearer):
+class _BaseJWTBearer(HTTPBearer):
     """
-    A FastAPI security dependency that validates JWT tokens for the Execution API.
+    Base class for JWT validation in the Execution API.
 
-    This validates tokens are signed and that the ``sub`` is a UUID. Queue-scoped tokens
-    (with scope="queue") are rejected - they can only be used on the /run endpoint.
-
-    The dependency result will be a `TIToken` object containing the ``id`` UUID (from the ``sub``)
-    and other validated claims.
+    Validates JWT tokens are properly signed and extracts claims. Subclasses
+    handle scope-specific validation.
     """
 
     def __init__(
@@ -88,14 +85,8 @@ async def __call__(  # type: ignore[override]
                 validators = self.required_claims
             claims = await validator.avalidated_claims(creds.credentials, validators)
 
-            # Reject queue-scoped tokens - they can only be used on /run endpoint
-            # Only check if scope claim is present (allows backwards compatibility with tests)
-            scope = claims.get("scope")
-            if scope is not None and scope == TOKEN_SCOPE_QUEUE:
-                raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
-                    detail="Queue tokens cannot access this endpoint. Use the token from /run response.",
-                )
+            # Let subclasses validate scope
+            self._check_scope(claims)
 
             return TIToken(id=claims["sub"], claims=claims)
         except HTTPException:
@@ -104,64 +95,51 @@ async def __call__(  # type: ignore[override]
             log.warning("Failed to validate JWT", exc_info=True)
             raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"Invalid auth token: {err}")
 
+    def _check_scope(self, claims: dict[str, Any]) -> None:
+        """Override in subclasses to validate scope. Raise HTTPException if invalid."""
+        pass
 
-class JWTBearerQueueScope(HTTPBearer):
-    """
-    JWT auth dependency that ONLY accepts queue-scoped tokens.
 
-    Used exclusively by the /run endpoint. Queue tokens have scope="queue" and are
-    long-lived to survive executor queue wait times. The /run endpoint validates
-    the queue token and issues a short-lived execution token for subsequent API calls.
+class JWTBearer(_BaseJWTBearer):
     """
+    JWT validation that rejects workload-scoped tokens.
 
-    def __init__(self, path_param_name: str | None = None):
-        super().__init__(auto_error=False)
-        self.path_param_name = path_param_name
+    Used for most Execution API endpoints. Workload-scoped tokens can only be used
+    on the /run endpoint, which exchanges them for short-lived execution tokens.
+    """
 
-    async def __call__(  # type: ignore[override]
-        self,
-        request: Request,
-        services=DepContainer,
-    ) -> TIToken | None:
-        creds = await super().__call__(request)
-        if not creds:
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing auth token")
+    def _check_scope(self, claims: dict[str, Any]) -> None:
+        if claims.get("scope") == TOKEN_SCOPE_WORKLOAD:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="Workload tokens cannot access this endpoint. Use the token from /run response.",
+            )
 
-        validator: JWTValidator = await services.aget(JWTValidator)
 
-        try:
-            if self.path_param_name:
-                id = request.path_params[self.path_param_name]
-                validators: dict[str, Any] = {"sub": {"essential": True, "value": id}}
-            else:
-                validators = {}
-            claims = await validator.avalidated_claims(creds.credentials, validators)
+class JWTBearerWorkloadScope(_BaseJWTBearer):
+    """
+    JWT validation that ONLY accepts workload-scoped tokens.
 
-            # Only accept queue-scoped tokens (if scope claim is present)
-            # This allows backwards compatibility with tests that don't set scope
-            scope = claims.get("scope")
-            if scope is not None and scope != TOKEN_SCOPE_QUEUE:
-                raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
-                    detail="This endpoint requires a queue-scoped token",
-                )
+    Used exclusively by the /run endpoint. Workload tokens have scope="ExecuteTaskWorkload"
+    and are long-lived to survive executor queue wait times. The /run endpoint validates
+    the workload token and issues a short-lived execution token for subsequent API calls.
+    """
 
-            return TIToken(id=claims["sub"], claims=claims)
-        except HTTPException:
-            raise
-        except Exception as err:
-            log.warning("Failed to validate JWT", exc_info=True)
-            raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"Invalid auth token: {err}")
+    def _check_scope(self, claims: dict[str, Any]) -> None:
+        scope = claims.get("scope")
+        # Reject if scope is explicitly set to something other than workload scope
+        if scope is not None and scope != TOKEN_SCOPE_WORKLOAD:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="This endpoint requires a workload-scoped token",
+            )
 
 
 JWTBearerDep: TIToken = Depends(JWTBearer())
 
 # This checks that the UUID in the url matches the one in the token for us.
 JWTBearerTIPathDep = Depends(JWTBearer(path_param_name="task_instance_id"))
 
-# For /run endpoint only - accepts queue-scoped tokens and validates task_instance_id
-JWTBearerQueueDep = Depends(JWTBearerQueueScope(path_param_name="task_instance_id"))
-
 
 async def get_team_name_dep(session: AsyncSessionDep, token=JWTBearerDep) -> str | None:
     """Return the team name associated to the task (if any)."""

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, Response, status
+from fastapi import Body, Depends, HTTPException, Query, Response, status
 from pydantic import JsonValue
 from sqlalchemy import func, or_, tuple_, update
 from sqlalchemy.engine import CursorResult, Row
@@ -60,7 +60,7 @@
     TISuccessStatePayload,
     TITerminalStatePayload,
 )
-from airflow.api_fastapi.execution_api.deps import DepContainer, JWTBearerQueueDep, JWTBearerTIPathDep
+from airflow.api_fastapi.execution_api.deps import DepContainer, JWTBearerTIPathDep, JWTBearerWorkloadScope
 from airflow.exceptions import TaskNotFound
 from airflow.models.asset import AssetActive
 from airflow.models.dag import DagModel
@@ -90,6 +90,9 @@
 
 log = structlog.get_logger(__name__)
 
+# For /run endpoint only - accepts workload-scoped tokens and validates task_instance_id
+JWTBearerWorkloadDep = Depends(JWTBearerWorkloadScope(path_param_name="task_instance_id"))
+
 
 @router.patch(
     "/{task_instance_id}/run",
@@ -100,7 +103,7 @@
         HTTP_422_UNPROCESSABLE_CONTENT: {"description": "Invalid payload for the state transition"},
     },
     response_model_exclude_unset=True,
-    dependencies=[JWTBearerQueueDep],
+    dependencies=[JWTBearerWorkloadDep],
 )
 async def ti_run(
     task_instance_id: UUID,

airflow-core/src/airflow/config_templates/config.yml

@@ -1822,14 +1822,14 @@ execution_api:
       type: integer
       example: ~
       default: "600"
-    jwt_queue_token_expiration_time:
+    jwt_workload_token_expiration_time:
       description: |
-        Number in seconds until the queue JWT token expires. Queue tokens are long-lived tokens
+        Number in seconds until the workload JWT token expires. Workload tokens are long-lived tokens
         sent with task workloads to executors (e.g., Celery). They can only be used to call
         the /run endpoint, which then issues a short-lived execution token.
 
         This should be set long enough to cover the maximum expected queue wait time.
-      version_added: 3.1.0
+      version_added: 3.1.7
       type: integer
       example: ~
       default: "86400"

airflow-core/src/airflow/executors/workloads.py

@@ -46,12 +46,12 @@ class BaseWorkload(BaseModel):
     @staticmethod
     def generate_token(sub_id: str, generator: JWTGenerator | None = None) -> str:
         """
-        Generate a queue-scoped token for this workload.
+        Generate a workload-scoped token for this workload.
 
-        Queue tokens are long-lived and can only be used on the /run endpoint,
+        Workload tokens are long-lived and can only be used on the /run endpoint,
         which exchanges them for short-lived execution tokens.
         """
-        return generator.generate_queue_token(sub_id) if generator else ""
+        return generator.generate_workload_token(sub_id) if generator else ""
 
 
 class BundleInfo(BaseModel):