Add token refresh mechanism for Execution API (#59553)

Tasks waiting in Celery queue may have their JWT tokens expire before
execution starts. This adds a token refresh endpoint that allows the
supervisor to refresh expired tokens before task execution.

Changes:
- Add /token/refresh endpoint to Execution API
- Add client-side token refresh logic in supervisor.py
- Add tests for the new endpoint

Fixes: #59553

Author: anishgirianish

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

@@ -29,13 +29,16 @@
     hitl,
     task_instances,
     task_reschedules,
+    tokens,
     variables,
     xcoms,
 )
 
 execution_api_router = APIRouter()
 execution_api_router.include_router(health.router, prefix="/health", tags=["Health"])
 
+execution_api_router.include_router(tokens.router, prefix="/token", tags=["Token"])
+
 # _Every_ single endpoint under here must be authenticated. Some do further checks on top of these
 authenticated_router = VersionedAPIRouter(dependencies=[JWTBearerDep])  # type: ignore[list-item]
 

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

@@ -0,0 +1,162 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Token refresh endpoint for the Execution API.
+
+This module provides an endpoint for workers to refresh expired JWT tokens.
+When a task waits in a queue (e.g., Celery) for longer than the token validity
+period, the worker can use this endpoint to obtain a fresh token before
+starting task execution.
+"""
+
+from __future__ import annotations
+
+import jwt
+import structlog
+from fastapi import APIRouter, Body, HTTPException, status
+from pydantic import BaseModel
+from sqlalchemy.exc import NoResultFound
+from sqlalchemy.sql import select
+
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.execution_api.app import _jwt_generator, _jwt_validator
+from airflow.models.taskinstance import TaskInstance as TI
+from airflow.utils.state import TaskInstanceState
+
+log = structlog.get_logger(logger_name=__name__)
+
+router = APIRouter()
+
+
+class TokenRefreshRequest(BaseModel):
+    """Request body for token refresh."""
+
+    token: str
+    """The expired (or about to expire) JWT token to refresh."""
+
+
+class TokenRefreshResponse(BaseModel):
+    """Response body for token refresh."""
+
+    access_token: str
+    """The new JWT token."""
+
+
+@router.post(
+    "/refresh",
+    status_code=status.HTTP_200_OK,
+    responses={
+        status.HTTP_400_BAD_REQUEST: {"description": "Invalid token format or signature"},
+        status.HTTP_403_FORBIDDEN: {"description": "Task is not in a valid state for token refresh"},
+        status.HTTP_404_NOT_FOUND: {"description": "Task instance not found"},
+    },
+)
+def refresh_token(
+    session: SessionDep,
+    body: TokenRefreshRequest = Body(...),
+) -> TokenRefreshResponse:
+    """
+    Refresh an expired or expiring JWT token for task execution.
+
+    This endpoint allows workers to obtain a fresh token when the original token
+    (embedded in the workload) has expired while waiting in a queue. The endpoint:
+
+    1. Validates the token signature (but ignores expiration)
+    2. Extracts the task instance ID from the token
+    3. Verifies the task is in QUEUED or RUNNING state
+    4. Issues a fresh token
+
+    This is necessary for distributed executors like Celery where tasks may wait
+    in queues longer than the token validity period.
+    """
+    token = body.token
+    validator = _jwt_validator()
+    generator = _jwt_generator()
+
+    try:
+        if validator.secret_key:
+            key = validator.secret_key
+        elif validator.jwks:
+            header = jwt.get_unverified_header(token)
+            kid = header.get("kid")
+            if not kid:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail="Token missing 'kid' header",
+                )
+            from asgiref.sync import async_to_sync
+
+            key = async_to_sync(validator.jwks.get_key)(kid)
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="No validation key configured",
+            )
+
+        claims = jwt.decode(
+            token,
+            key,
+            audience=validator.audience,
+            issuer=validator.issuer,
+            options={
+                "verify_exp": False,
+                "require": ["sub"],
+            },
+            algorithms=validator.algorithm,
+            leeway=validator.leeway,
+        )
+    except jwt.InvalidTokenError as e:
+        log.warning("Invalid token for refresh", error=str(e))
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Invalid token: {e}",
+        )
+
+    task_instance_id = claims.get("sub")
+    if not task_instance_id:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Token missing 'sub' claim (task instance ID)",
+        )
+
+    try:
+        ti_state = session.execute(select(TI.state).where(TI.id == task_instance_id)).scalar_one()
+    except NoResultFound:
+        log.warning("Task instance not found for token refresh", task_instance_id=task_instance_id)
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="Task instance not found",
+        )
+
+    valid_states = {TaskInstanceState.QUEUED, TaskInstanceState.RUNNING}
+    if ti_state not in valid_states:
+        log.warning(
+            "Task not in valid state for token refresh",
+            task_instance_id=task_instance_id,
+            state=ti_state,
+        )
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=f"Task is in '{ti_state}' state, token refresh only allowed for QUEUED or RUNNING tasks",
+        )
+
+    new_token = generator.generate({"sub": task_instance_id})
+
+    log.info("Token refreshed successfully", task_instance_id=task_instance_id)
+
+    return TokenRefreshResponse(access_token=new_token)