Add single-use callback token for deadline callback context fetch (#71192)

Deadline callbacks run in a subprocess that needs to read the DagRun
context (and connections/variables/xcoms) from the Execution API. PR
also accept the long-lived ``workload`` token via ``token:workload``
opt-ins. That over-broadened the workload token's reach (scope creep):
a long-lived token could read arbitrary DagRun/connection/variable/xcom
data for the whole queue-wait lifetime, and Ash asked for a single-use
credential instead.

This re-lands the security core of #66608 with a tighter design:

---------

Co-authored-by: Ash Berlin-Taylor <ash_github@firemirror.com>

Author: seanghaeli

airflow-core/newsfragments/71192.significant.rst

@@ -0,0 +1,12 @@
+Callbacks now redeem a single-use token before they run
+
+A worker now redeems a callback's single-use token with the API server
+(``PATCH /execution/callbacks/{callback_id}/run``) before importing and running the callback.
+Redeeming the token moves the callback from queued to running and swaps it for a short-lived
+execution-scoped token. Before this change a callback ran without the worker making any
+authenticated call, so the token minted for it was never checked.
+
+Because the token is single-use, a redelivered or replayed message that reaches a callback that is
+already running or finished is refused rather than run again. A worker on 3.4 needs an API server
+that serves this endpoint; a worker talking to an older API server negotiates the API version and
+behaves as though the endpoint is not there.

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=["cb: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.db.dags import eager_load_teams
@@ -75,6 +74,7 @@
     CurrentTIToken,
     ExecutionAPIRoute,
     get_team_name_for_ti,
+    issue_execution_token,
     require_auth,
 )
 from airflow.api_fastapi.execution_api.services.task_instances import (
@@ -346,9 +346,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,15 @@
 from typing import Any, get_args
 
 import structlog
-from fastapi import Depends, HTTPException, Request, status
+import svcs
+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 +197,26 @@ async def require_auth(
                 status_code=status.HTTP_403_FORBIDDEN,
                 detail="Token subject does not match connection test ID",
             )
+    elif "cb: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: svcs.Container, 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/api_fastapi/execution_api/versions/__init__.py

@@ -51,11 +51,14 @@
     AddTeamNameField,
     AddVariableKeysEndpoint,
 )
-from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext
+from airflow.api_fastapi.execution_api.versions.v2026_10_30 import (
+    AddArgBindingsToTIRunContext,
+    AddCallbackRunEndpoint,
+)
 
 bundle = VersionBundle(
     HeadVersion(),
-    Version("2026-10-30", AddArgBindingsToTIRunContext),
+    Version("2026-10-30", AddArgBindingsToTIRunContext, AddCallbackRunEndpoint),
     Version(
         "2026-06-30",
         AddVariableKeysEndpoint,

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py

@@ -19,8 +19,10 @@
 
 from cadwyn import (
     ResponseInfo,
+    VersionChange,
     VersionChangeWithSideEffects,
     convert_response_to_previous_version_for,
+    endpoint,
     schema,
 )
 
@@ -40,3 +42,13 @@ class AddArgBindingsToTIRunContext(VersionChangeWithSideEffects):
     def remove_arg_bindings_field(response: ResponseInfo) -> None:  # type: ignore[misc]
         """Strip ``arg_bindings`` from the run context for older clients."""
         response.body.pop("arg_bindings", None)
+
+
+class AddCallbackRunEndpoint(VersionChange):
+    """Add the callbacks/{callback_id}/run endpoint a worker uses to exchange its single-use callback token."""
+
+    description = __doc__
+
+    instructions_to_migrate_to_previous_version = (
+        endpoint("/callbacks/{callback_id}/run", ["PATCH"]).didnt_exist,
+    )

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

@@ -21,7 +21,7 @@
 import os
 from abc import ABC, abstractmethod
 from collections.abc import Hashable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -83,13 +83,16 @@ class BaseWorkloadSchema(BaseModel):
     token: str = Field(repr=False)
     """The identity token for this workload"""
 
-    @staticmethod
-    def generate_token(sub_id: str, generator: JWTGenerator | None = None) -> str:
+    token_scope: ClassVar[str] = "workload"
+    """Scope claim stamped into tokens minted for this workload type."""
+
+    @classmethod
+    def generate_token(cls, sub_id: str, generator: JWTGenerator | None = None) -> 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": cls.token_scope},
             valid_for=valid_for,
         )