Add single-use callback token for deadline callback context fetch

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

@@ -148,10 +148,10 @@ 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":
+                    # Workload and callback tokens are long-lived and meant to survive
+                    # queue wait times so avoid refreshing them. If avalidated_claims
+                    # raises for such a token, the outer except handles it.
+                    if claims.get("scope") in ("workload", "callback"):
                         return response
 
                     now = int(time.time())

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/token.py

@@ -24,7 +24,7 @@
 
 from airflow.api_fastapi.core_api.base import BaseModel
 
-TokenScope = Literal["execution", "workload"]
+TokenScope = Literal["execution", "workload", "callback"]
 
 
 class TIClaims(BaseModel):

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

@@ -23,6 +23,7 @@
     asset_events,
     asset_state_store,
     assets,
+    callbacks,
     connection_tests,
     connections,
     dag_runs,
@@ -53,6 +54,7 @@
     connection_tests.router, prefix="/connection-tests", tags=["Connection Tests"]
 )
 authenticated_router.include_router(connections.router, prefix="/connections", tags=["Connections"])
+authenticated_router.include_router(callbacks.router, prefix="/callbacks", tags=["Callbacks"])
 authenticated_router.include_router(dag_runs.router, prefix="/dag-runs", tags=["Dag Runs"])
 authenticated_router.include_router(dags.router, prefix="/dags", tags=["Dags"])
 authenticated_router.include_router(task_instances.router, prefix="/task-instances", tags=["Task Instances"])

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

@@ -0,0 +1,85 @@
+# 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.
+from __future__ import annotations
+
+from uuid import UUID
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Response, Security, status
+
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.api_fastapi.execution_api.deps import DepContainer
+from airflow.api_fastapi.execution_api.security import (
+    ExecutionAPIRoute,
+    issue_execution_token,
+    require_auth,
+)
+from airflow.models.callback import Callback
+from airflow.utils.state import CallbackState
+
+router = VersionedAPIRouter(
+    route_class=ExecutionAPIRoute,
+    dependencies=[
+        Security(require_auth, scopes=["callback:self", "token:callback"]),
+    ],
+)
+
+
+@router.patch(
+    "/{callback_id}/run",
+    status_code=status.HTTP_204_NO_CONTENT,
+    responses=create_openapi_http_exception_doc(
+        [
+            (status.HTTP_404_NOT_FOUND, "Callback not found"),
+            (status.HTTP_409_CONFLICT, "The callback token was already exchanged"),
+        ]
+    ),
+)
+def run_callback(
+    callback_id: UUID,
+    response: Response,
+    session: SessionDep,
+    services=DepContainer,
+) -> None:
+    """Exchange a single-use callback token for a short-lived execution token."""
+    callback = session.get(Callback, callback_id, with_for_update=True)
+    if callback is None:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail={
+                "reason": "not_found",
+                "message": f"Callback {callback_id} not found",
+            },
+        )
+
+    if callback.state != CallbackState.QUEUED:
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail={
+                "reason": "invalid_state",
+                "message": (
+                    f"Callback {callback_id} is in state {callback.state}; its token can only be "
+                    "exchanged once while QUEUED."
+                ),
+                "previous_state": callback.state,
+            },
+        )
+
+    callback.state = CallbackState.RUNNING
+
+    issue_execution_token(services, response, sub=str(callback_id))

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

@@ -43,7 +43,6 @@
 from airflow._shared.observability.traces import override_ids
 from airflow._shared.state import TaskScope
 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
@@ -73,6 +72,7 @@
     CurrentTIToken,
     ExecutionAPIRoute,
     get_team_name_for_ti,
+    issue_execution_token,
     require_auth,
 )
 from airflow.configuration import conf
@@ -324,9 +324,7 @@ def ti_run(
 
     # 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
+        issue_execution_token(services, response, sub=str(task_instance_id))
 
     return context
 

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

@@ -70,14 +70,14 @@
 from typing import Any, get_args
 
 import structlog
-from fastapi import Depends, HTTPException, Request, status
+from fastapi import Depends, HTTPException, Request, Response, status
 from fastapi.params import Security as SecurityParam
 from fastapi.routing import APIRoute
 from fastapi.security import HTTPBearer, SecurityScopes
 from pydantic import ValidationError
 from sqlalchemy import select
 
-from airflow.api_fastapi.auth.tokens import JWTValidator
+from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
 from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken, TokenScope
 from airflow.api_fastapi.execution_api.deps import DepContainer
 
@@ -196,13 +196,26 @@ async def require_auth(
                 status_code=status.HTTP_403_FORBIDDEN,
                 detail="Token subject does not match connection test ID",
             )
+    elif "callback:self" in security_scopes.scopes:
+        cb_self_id = str(request.path_params["callback_id"])
+        if str(token.id) != cb_self_id:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="Token subject does not match callback ID",
+            )
 
     return token
 
 
 CurrentTIToken: TIToken = Depends(require_auth)
 
 
+def issue_execution_token(services: Any, response: Response, sub: str) -> None:
+    """Mint an ``execution``-scoped token and set it on the ``Refreshed-API-Token`` header."""
+    generator: JWTGenerator = services.get(JWTGenerator)
+    response.headers["Refreshed-API-Token"] = generator.generate(extras={"sub": sub, "scope": "execution"})
+
+
 class ExecutionAPIRoute(APIRoute):
     """
     Custom route class that precomputes allowed token types from Security scopes.

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

@@ -84,12 +84,12 @@ class BaseWorkloadSchema(BaseModel):
     """The identity token for this workload"""
 
     @staticmethod
-    def generate_token(sub_id: str, generator: JWTGenerator | None = None) -> str:
+    def generate_token(sub_id: str, generator: JWTGenerator | None = None, scope: str = "workload") -> str:
         if not generator:
             return ""
         valid_for = conf.getfloat("scheduler", "task_queued_timeout")
         return generator.generate(
-            extras={"sub": sub_id, "scope": "workload"},
+            extras={"sub": sub_id, "scope": scope},
             valid_for=valid_for,
         )
 

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

@@ -127,7 +127,7 @@ def make(
         return cls(
             callback=CallbackDTO.model_validate(callback, from_attributes=True),
             dag_rel_path=dag_rel_path or Path(dag_run.dag_model.relative_fileloc or ""),
-            token=cls.generate_token(str(callback.id), generator),
+            token=cls.generate_token(str(callback.id), generator, scope="callback"),
             log_path=fname,
             bundle_info=bundle_info,
         )

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

@@ -58,7 +58,10 @@ async def mock_require_auth(request: Request) -> TIToken:
 
         raw_id = request.path_params.get(
             "task_instance_id",
-            request.path_params.get("connection_test_id", "00000000-0000-0000-0000-000000000000"),
+            request.path_params.get(
+                "connection_test_id",
+                request.path_params.get("callback_id", "00000000-0000-0000-0000-000000000000"),
+            ),
         )
         try:
             ti_id = UUID(raw_id)

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

@@ -80,6 +80,16 @@ def test_ignores_non_token_scopes(self):
         )
         assert route.allowed_token_types == frozenset({"execution"})
 
+    def test_extracts_callback_token_scope(self):
+        route = ExecutionAPIRoute(
+            path="/test",
+            endpoint=lambda: None,
+            dependencies=[
+                Security(require_auth, scopes=["callback:self", "token:callback"]),
+            ],
+        )
+        assert route.allowed_token_types == frozenset({"callback"})
+
     def test_rejects_invalid_token_types(self):
         with pytest.raises(ValueError, match="Invalid token types"):
             ExecutionAPIRoute(

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

@@ -49,6 +49,8 @@
     # Connection test routes run from a queued worker context (workload-only).
     "PATCH /connection-tests/{connection_test_id}": {"workload"},
     "GET /connection-tests/{connection_test_id}/connection": {"workload"},
+    # Callback /run exchanges a single-use callback token for an execution token.
+    "PATCH /callbacks/{callback_id}/run": {"callback"},
 }