From 02af3fdb02d2b82d4419a672b56cb54af42ad73d Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Mon, 13 Jul 2026 22:34:57 +0000 Subject: [PATCH 1/4] Add single-use callback token for deadline callback context fetch 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: * New ``callback`` token scope. The executor mints a ``callback``-scoped token whose ``sub`` is the callback id. That token is accepted ONLY on the new ``PATCH /callbacks/{id}/run`` exchange endpoint (pinned by a ``callback:self`` scope) and nowhere else. * Single-use exchange. ``run_callback`` performs an atomic ``QUEUED -> RUNNING`` transition under ``SELECT ... FOR UPDATE`` and returns a fresh short-lived ``execution`` token via the ``Refreshed-API-Token`` header. A second presentation of the same callback token finds the row already RUNNING and gets 409, so the token is exchangeable exactly once. This mirrors the Task Instance ``/run`` pattern. The subprocess calls the exchange once, before any context read, then uses the execution token for the actual reads. * Reverted the #66608 ``token:workload`` opt-ins on the DagRun, connections, variables, and xcoms GET endpoints back to execution-only. * ``JWTReissueMiddleware`` skips ``callback`` tokens (as it already does for ``workload``), since they are long-lived and exchanged exactly once. Two shared-helper refactors keep the new code DRY and also clean up pre-existing duplication: * ``require_auth`` now drives the three ``*:self`` scope checks (``ti:self`` / ``ct:self`` / ``callback:self``) from a single ``SELF_SCOPE_PATH_PARAMS`` mapping + loop instead of three near-identical branches. Behavior for ``ti:self`` and ``ct:self`` is unchanged. * New ``issue_execution_token`` helper mints an execution-scope token and sets the ``Refreshed-API-Token`` header. It replaces the duplicated mint block in the Task Instance ``/run`` handler (still guarded by the ``workload`` scope check) and the new callback exchange endpoint. --- .../airflow/api_fastapi/execution_api/app.py | 6 +- .../execution_api/datamodels/token.py | 2 +- .../execution_api/routes/__init__.py | 2 + .../execution_api/routes/callbacks.py | 104 ++++++++++ .../execution_api/routes/task_instances.py | 6 +- .../api_fastapi/execution_api/security.py | 49 +++-- .../src/airflow/executors/workloads/base.py | 4 +- .../airflow/executors/workloads/callback.py | 2 +- .../api_fastapi/execution_api/conftest.py | 5 +- .../execution_api/test_security.py | 10 + .../versions/head/test_callbacks.py | 178 ++++++++++++++++++ task-sdk/src/airflow/sdk/api/client.py | 26 +++ .../sdk/execution_time/callback_supervisor.py | 4 + task-sdk/tests/task_sdk/api/test_client.py | 36 ++++ 14 files changed, 405 insertions(+), 29 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 5a87cdb81e04c..ea0a55aa9f395 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -148,10 +148,8 @@ 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 (exchanged once, not refreshed here). + if claims.get("scope") in ("workload", "callback"): return response now = int(time.time()) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/token.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/token.py index 4c3b935f5aa62..eb858fac1dfc4 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/token.py +++ b/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): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py index 7b19f3ddd3055..5662378c87087 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/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"]) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py new file mode 100644 index 0000000000000..ce5fcde289087 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py @@ -0,0 +1,104 @@ +# 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 + +import structlog +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.datamodels.token import TIToken +from airflow.api_fastapi.execution_api.deps import DepContainer +from airflow.api_fastapi.execution_api.security import ( + CurrentTIToken, + ExecutionAPIRoute, + issue_execution_token, + require_auth, +) +from airflow.models.callback import Callback +from airflow.utils.state import CallbackState + +log = structlog.get_logger(logger_name=__name__) + +# The ``callback``-scoped token is accepted only on this router; ``callback:self`` pins it +# to its own callback id and it can reach only the exchange endpoint below. +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, + token: TIToken = CurrentTIToken, +) -> None: + """ + Exchange a single-use callback token for a short-lived execution token. + + The deadline-callback subprocess calls this once, before reading any + context. The exchange is gated on an atomic ``QUEUED -> RUNNING`` transition + under ``SELECT ... FOR UPDATE``: the first call transitions the row and + returns a fresh ``execution`` token via the ``Refreshed-API-Token`` header; + any second presentation of the same callback token finds the row already in + ``RUNNING`` (or terminal) and gets ``409``. This makes the callback token + exchangeable exactly once, mirroring the Task Instance ``/run`` pattern. + """ + 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": str(callback.state) if callback.state is not None else None, + }, + ) + + callback.state = CallbackState.RUNNING + log.info("Callback token exchanged; marked running", callback_id=str(callback_id)) + + issue_execution_token(services, response, sub=str(callback_id)) 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 index c1bac7960234d..b32831527dad6 100644 --- 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 @@ -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 @@ -326,9 +326,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 diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/security.py b/airflow-core/src/airflow/api_fastapi/execution_api/security.py index 9de3493061f32..ae372e5678b50 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/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 @@ -85,6 +85,14 @@ VALID_TOKEN_TYPES: frozenset[str] = frozenset(get_args(TokenScope)) +# ``*:self`` Security scopes pin a token to a single resource. Each maps the scope name to +# the route path parameter carrying the resource id that the JWT ``sub`` must match. +SELF_SCOPE_PATH_PARAMS: dict[str, str] = { + "ti:self": "task_instance_id", + "ct:self": "connection_test_id", + "callback:self": "callback_id", +} + _REQUEST_SCOPE_TOKEN_KEY = "ti_token" @@ -182,20 +190,15 @@ async def require_auth( f"Allowed types: {', '.join(sorted(allowed_token_types))}", ) - if "ti:self" in security_scopes.scopes: - ti_self_id = str(request.path_params["task_instance_id"]) - if str(token.id) != ti_self_id: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="Token subject does not match task instance ID", - ) - elif "ct:self" in security_scopes.scopes: - ct_self_id = str(request.path_params["connection_test_id"]) - if str(token.id) != ct_self_id: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="Token subject does not match connection test ID", - ) + # Enforce the first ``*:self`` scope present: the JWT ``sub`` must match the path param. + for scope, path_param in SELF_SCOPE_PATH_PARAMS.items(): + if scope in security_scopes.scopes: + if str(token.id) != str(request.path_params[path_param]): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Token subject does not match {path_param}", + ) + break return token @@ -203,6 +206,20 @@ async def require_auth( CurrentTIToken: TIToken = Depends(require_auth) +def issue_execution_token(services: Any, response: Response, sub: str) -> None: + """ + Mint a short-lived ``execution``-scoped token and set it on the response. + + Used by endpoints that swap a longer-lived token (``workload`` on the Task Instance + ``/run`` endpoint, ``callback`` on the callback exchange endpoint) for an ``execution`` + token: it mints via the request-scoped ``JWTGenerator`` and sets the ``Refreshed-API-Token`` + header the SDK client adopts automatically. ``JWTReissueMiddleware`` skips those longer-lived + scopes, so this is the single place the swap happens. + """ + 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. diff --git a/airflow-core/src/airflow/executors/workloads/base.py b/airflow-core/src/airflow/executors/workloads/base.py index 41334d68f3038..98788bc49e59b 100644 --- a/airflow-core/src/airflow/executors/workloads/base.py +++ b/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, ) diff --git a/airflow-core/src/airflow/executors/workloads/callback.py b/airflow-core/src/airflow/executors/workloads/callback.py index c1842a5900685..702e849fa407b 100644 --- a/airflow-core/src/airflow/executors/workloads/callback.py +++ b/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, ) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py b/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py index fda87d21e21a9..be6e7f259c54d 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/conftest.py +++ b/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) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py index e98d918385933..7fd2fcdb4f064 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py +++ b/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( diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py new file mode 100644 index 0000000000000..d39de34649a51 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py @@ -0,0 +1,178 @@ +# 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 unittest import mock + +import jwt +import pytest + +from airflow.api_fastapi.auth.tokens import JWTValidator +from airflow.api_fastapi.execution_api.app import lifespan +from airflow.api_fastapi.execution_api.security import require_auth +from airflow.models.callback import CallbackFetchMethod, CallbackState, ExecutorCallback +from airflow.sdk.definitions.callback import SyncCallback + +from tests_common.test_utils.db import clear_db_callbacks + +pytestmark = pytest.mark.db_test + + +def sync_callback(): + """Empty (sync) callable used for unit tests.""" + + +def _make_queued_callback(session, dag_id="test_dag"): + """Create and persist an ExecutorCallback in QUEUED state.""" + callback = ExecutorCallback( + callback_def=SyncCallback(sync_callback, kwargs={}), + fetch_method=CallbackFetchMethod.IMPORT_PATH, + ) + callback.data["dag_id"] = dag_id + callback.state = CallbackState.QUEUED + session.add(callback) + session.commit() + return callback + + +class TestRunCallback: + @pytest.fixture(autouse=True) + def setup_teardown(self): + clear_db_callbacks() + yield + clear_db_callbacks() + + def test_run_marks_running_and_swaps_token(self, client, session): + """First call transitions QUEUED -> RUNNING and returns a fresh execution token.""" + callback = _make_queued_callback(session) + + response = client.patch(f"/execution/callbacks/{callback.id}/run") + + assert response.status_code == 204 + + # A real execution-scope JWT is minted with the callback id as its subject. + refreshed = response.headers["Refreshed-API-Token"] + payload = jwt.decode(refreshed, options={"verify_signature": False}) + assert payload["scope"] == "execution" + assert payload["sub"] == str(callback.id) + + session.expire_all() + session.refresh(callback) + assert callback.state == CallbackState.RUNNING + + def test_run_is_single_use(self, client, session): + """A second exchange of the same callback token is rejected with 409.""" + callback = _make_queued_callback(session) + + first = client.patch(f"/execution/callbacks/{callback.id}/run") + assert first.status_code == 204 + + second = client.patch(f"/execution/callbacks/{callback.id}/run") + assert second.status_code == 409 + detail = second.json()["detail"] + assert detail["reason"] == "invalid_state" + assert detail["previous_state"] == CallbackState.RUNNING + + @pytest.mark.parametrize( + "state", + [CallbackState.RUNNING, CallbackState.SUCCESS, CallbackState.FAILED, CallbackState.PENDING], + ) + def test_run_rejects_non_queued_state(self, client, session, state): + """The token can only be exchanged while the callback is QUEUED.""" + callback = _make_queued_callback(session) + callback.state = state + session.commit() + + response = client.patch(f"/execution/callbacks/{callback.id}/run") + assert response.status_code == 409 + assert response.json()["detail"]["reason"] == "invalid_state" + + def test_run_returns_404_for_nonexistent(self, client): + """Exchanging a token for an unknown callback returns 404.""" + response = client.patch("/execution/callbacks/00000000-0000-0000-0000-000000000000/run") + assert response.status_code == 404 + assert response.json()["detail"]["reason"] == "not_found" + + +@pytest.fixture +def _use_real_jwt_bearer(exec_app): + """Remove the mock require_auth override so the real scope validation runs.""" + exec_app.dependency_overrides.pop(require_auth, None) + + +@pytest.mark.usefixtures("_use_real_jwt_bearer") +def test_run_accepts_callback_token_matching_sub(client, session): + """The exchange endpoint accepts a callback-scope JWT whose sub is the callback id.""" + clear_db_callbacks() + callback = _make_queued_callback(session) + + validator = mock.AsyncMock(spec=JWTValidator) + validator.avalidated_claims.return_value = { + "sub": str(callback.id), + "scope": "callback", + "exp": 9999999999, + "iat": 1000000000, + "nbf": 1000000000, + } + lifespan.registry.register_value(JWTValidator, validator) + + resp = client.patch(f"/execution/callbacks/{callback.id}/run") + assert resp.status_code == 204, resp.json() + clear_db_callbacks() + + +@pytest.mark.usefixtures("_use_real_jwt_bearer") +def test_run_rejects_callback_token_for_other_callback(client, session): + """callback:self blocks a callback token whose sub is a different callback id.""" + clear_db_callbacks() + callback = _make_queued_callback(session) + + validator = mock.AsyncMock(spec=JWTValidator) + validator.avalidated_claims.return_value = { + # sub belongs to a different callback than the one in the URL. + "sub": "11111111-1111-1111-1111-111111111111", + "scope": "callback", + "exp": 9999999999, + "iat": 1000000000, + "nbf": 1000000000, + } + lifespan.registry.register_value(JWTValidator, validator) + + resp = client.patch(f"/execution/callbacks/{callback.id}/run") + assert resp.status_code == 403, resp.json() + clear_db_callbacks() + + +@pytest.mark.usefixtures("_use_real_jwt_bearer") +def test_run_rejects_execution_scope_token(client, session): + """Only callback-scope tokens are accepted on the exchange endpoint.""" + clear_db_callbacks() + callback = _make_queued_callback(session) + + validator = mock.AsyncMock(spec=JWTValidator) + validator.avalidated_claims.return_value = { + "sub": str(callback.id), + "scope": "execution", + "exp": 9999999999, + "iat": 1000000000, + "nbf": 1000000000, + } + lifespan.registry.register_value(JWTValidator, validator) + + resp = client.patch(f"/execution/callbacks/{callback.id}/run") + assert resp.status_code == 403, resp.json() + clear_db_callbacks() diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 74a680053711c..850af57cfe270 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1085,6 +1085,26 @@ def update_state( self.client.patch(f"connection-tests/{id}", content=body.model_dump_json()) +class CallbackOperations: + __slots__ = ("client",) + + def __init__(self, client: Client): + self.client = client + + def run(self, callback_id: uuid.UUID) -> None: + """ + Exchange the single-use callback token for a short-lived execution token. + + Marks the callback RUNNING on the server and swaps this client's + ``callback``-scoped token for an ``execution`` token via the + ``Refreshed-API-Token`` response header (adopted automatically by + ``_update_auth``). Must be called once, before any context read; a + second call raises ``ServerResponseError`` (409) because the callback + is no longer QUEUED. + """ + self.client.patch(f"callbacks/{callback_id}/run") + + class BearerAuth(httpx.Auth): def __init__(self, token: str): self.token: str = token @@ -1292,6 +1312,12 @@ def connection_tests(self) -> ConnectionTestOperations: """Operations related to Connection Tests.""" return ConnectionTestOperations(self) + @lru_cache() # type: ignore[misc] + @property + def callbacks(self) -> CallbackOperations: + """Operations related to Callbacks.""" + return CallbackOperations(self) + @lru_cache() # type: ignore[misc] @property def dags(self) -> DagsOperations: diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py index 9830771701293..da9802d086861 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -433,6 +433,10 @@ def supervise_callback( else: logger = structlog.get_logger(logger_name="callback").bind() + # Swap the single-use callback token for an execution token before any context read. + # A duplicate delivery finds the callback already RUNNING and raises, so it runs at most once. + client.callbacks.run(UUID(id)) + try: process = CallbackSubprocess.start( id=id, diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 205b12bd4f2d8..57e2fb6b6e3d3 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -2093,3 +2093,39 @@ def handle_request(request: httpx.Request) -> httpx.Response: client = make_client(transport=httpx.MockTransport(handle_request)) result = client.asset_state_store.clear(uri="s3://bucket/key") assert result == OKResponse(ok=True) + + +class TestCallbackOperations: + def test_run_exchanges_token(self): + """run() PATCHes the exchange endpoint and adopts the refreshed execution token.""" + callback_id = uuid7() + + def handle_request(request: httpx.Request) -> httpx.Response: + assert request.method == "PATCH" + assert request.url.path == f"/callbacks/{callback_id}/run" + return httpx.Response( + status_code=204, + headers={"Refreshed-API-Token": "execution-token"}, + ) + + client = make_client(transport=httpx.MockTransport(handle_request)) + client.callbacks.run(callback_id) + + # The single-use callback token has been swapped for the execution token. + assert client.auth is not None + assert client.auth.token == "execution-token" + + def test_run_raises_on_conflict(self): + """A second exchange (callback no longer QUEUED) surfaces as ServerResponseError.""" + callback_id = uuid7() + + def handle_request(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=409, + json={"detail": {"reason": "invalid_state", "previous_state": "running"}}, + ) + + client = make_client(transport=httpx.MockTransport(handle_request)) + with pytest.raises(ServerResponseError) as err: + client.callbacks.run(callback_id) + assert err.value.response.status_code == 409 From 284cea5116db8500ffd517bde933fe4b16002bd7 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Wed, 15 Jul 2026 00:09:37 +0000 Subject: [PATCH 2/4] Declare callback /run route in NON_DEFAULT_TOKEN_POLICY The new PATCH /callbacks/{callback_id}/run route accepts the single-use 'callback' token, a non-default (non-execution) scope. The token scope boundary regression guard requires such routes to be declared explicitly in NON_DEFAULT_TOKEN_POLICY; add it there per the test's documented maintenance step. --- .../api_fastapi/execution_api/test_token_scope_boundaries.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py index bbf2be8704f46..e15e53ff1cab0 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py @@ -49,6 +49,9 @@ # 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"}, + # The callback /run endpoint exchanges a single-use callback token for a + # short-lived execution token (callback-only). + "PATCH /callbacks/{callback_id}/run": {"callback"}, } From bee9b82f3bab4bcd8da591925e39590390b57ec2 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Fri, 24 Jul 2026 23:40:26 +0000 Subject: [PATCH 3/4] Trim over-explained comments --- .../api_fastapi/execution_api/routes/callbacks.py | 11 ++--------- .../airflow/api_fastapi/execution_api/security.py | 13 ++++--------- .../execution_api/test_token_scope_boundaries.py | 3 +-- .../execution_api/versions/head/test_callbacks.py | 2 -- task-sdk/src/airflow/sdk/api/client.py | 12 ++++-------- .../sdk/execution_time/callback_supervisor.py | 3 +-- task-sdk/tests/task_sdk/api/test_client.py | 1 - 7 files changed, 12 insertions(+), 33 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py index ce5fcde289087..031ef87af75cc 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py @@ -37,8 +37,6 @@ log = structlog.get_logger(logger_name=__name__) -# The ``callback``-scoped token is accepted only on this router; ``callback:self`` pins it -# to its own callback id and it can reach only the exchange endpoint below. router = VersionedAPIRouter( route_class=ExecutionAPIRoute, dependencies=[ @@ -67,13 +65,8 @@ def run_callback( """ Exchange a single-use callback token for a short-lived execution token. - The deadline-callback subprocess calls this once, before reading any - context. The exchange is gated on an atomic ``QUEUED -> RUNNING`` transition - under ``SELECT ... FOR UPDATE``: the first call transitions the row and - returns a fresh ``execution`` token via the ``Refreshed-API-Token`` header; - any second presentation of the same callback token finds the row already in - ``RUNNING`` (or terminal) and gets ``409``. This makes the callback token - exchangeable exactly once, mirroring the Task Instance ``/run`` pattern. + Gated on an atomic ``QUEUED -> RUNNING`` transition, so the exchange + succeeds at most once; the new token is returned via ``Refreshed-API-Token``. """ callback = session.get(Callback, callback_id, with_for_update=True) if callback is None: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/security.py b/airflow-core/src/airflow/api_fastapi/execution_api/security.py index ae372e5678b50..b7bb6d9e153ab 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -85,8 +85,7 @@ VALID_TOKEN_TYPES: frozenset[str] = frozenset(get_args(TokenScope)) -# ``*:self`` Security scopes pin a token to a single resource. Each maps the scope name to -# the route path parameter carrying the resource id that the JWT ``sub`` must match. +# Maps each ``*:self`` scope to the path param whose value the JWT ``sub`` must match. SELF_SCOPE_PATH_PARAMS: dict[str, str] = { "ti:self": "task_instance_id", "ct:self": "connection_test_id", @@ -190,7 +189,6 @@ async def require_auth( f"Allowed types: {', '.join(sorted(allowed_token_types))}", ) - # Enforce the first ``*:self`` scope present: the JWT ``sub`` must match the path param. for scope, path_param in SELF_SCOPE_PATH_PARAMS.items(): if scope in security_scopes.scopes: if str(token.id) != str(request.path_params[path_param]): @@ -208,13 +206,10 @@ async def require_auth( def issue_execution_token(services: Any, response: Response, sub: str) -> None: """ - Mint a short-lived ``execution``-scoped token and set it on the response. + Mint an ``execution``-scoped token and set it on the ``Refreshed-API-Token`` header. - Used by endpoints that swap a longer-lived token (``workload`` on the Task Instance - ``/run`` endpoint, ``callback`` on the callback exchange endpoint) for an ``execution`` - token: it mints via the request-scoped ``JWTGenerator`` and sets the ``Refreshed-API-Token`` - header the SDK client adopts automatically. ``JWTReissueMiddleware`` skips those longer-lived - scopes, so this is the single place the swap happens. + ``JWTReissueMiddleware`` skips ``workload``/``callback`` tokens, so endpoints that + swap them for an execution token must set the header here themselves. """ generator: JWTGenerator = services.get(JWTGenerator) response.headers["Refreshed-API-Token"] = generator.generate(extras={"sub": sub, "scope": "execution"}) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py index e15e53ff1cab0..195503f2bf080 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_token_scope_boundaries.py @@ -49,8 +49,7 @@ # 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"}, - # The callback /run endpoint exchanges a single-use callback token for a - # short-lived execution token (callback-only). + # Callback /run exchanges a single-use callback token for an execution token. "PATCH /callbacks/{callback_id}/run": {"callback"}, } diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py index d39de34649a51..7c95df4641a19 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py @@ -64,7 +64,6 @@ def test_run_marks_running_and_swaps_token(self, client, session): assert response.status_code == 204 - # A real execution-scope JWT is minted with the callback id as its subject. refreshed = response.headers["Refreshed-API-Token"] payload = jwt.decode(refreshed, options={"verify_signature": False}) assert payload["scope"] == "execution" @@ -143,7 +142,6 @@ def test_run_rejects_callback_token_for_other_callback(client, session): validator = mock.AsyncMock(spec=JWTValidator) validator.avalidated_claims.return_value = { - # sub belongs to a different callback than the one in the URL. "sub": "11111111-1111-1111-1111-111111111111", "scope": "callback", "exp": 9999999999, diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index e0bea96e7f3ae..ec6aebf582afe 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1099,14 +1099,10 @@ def __init__(self, client: Client): def run(self, callback_id: uuid.UUID) -> None: """ - Exchange the single-use callback token for a short-lived execution token. - - Marks the callback RUNNING on the server and swaps this client's - ``callback``-scoped token for an ``execution`` token via the - ``Refreshed-API-Token`` response header (adopted automatically by - ``_update_auth``). Must be called once, before any context read; a - second call raises ``ServerResponseError`` (409) because the callback - is no longer QUEUED. + Exchange the single-use callback token for an execution token. + + Must be called before any context read; a second call raises + ``ServerResponseError`` (409). """ self.client.patch(f"callbacks/{callback_id}/run") diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py index da9802d086861..0b51a5a31580a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -433,8 +433,7 @@ def supervise_callback( else: logger = structlog.get_logger(logger_name="callback").bind() - # Swap the single-use callback token for an execution token before any context read. - # A duplicate delivery finds the callback already RUNNING and raises, so it runs at most once. + # Exchange the single-use callback token before any context read; raises on duplicate delivery. client.callbacks.run(UUID(id)) try: diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index cf5fe69c0edfd..a1274ec08f6d4 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -2200,7 +2200,6 @@ def handle_request(request: httpx.Request) -> httpx.Response: client = make_client(transport=httpx.MockTransport(handle_request)) client.callbacks.run(callback_id) - # The single-use callback token has been swapped for the execution token. assert client.auth is not None assert client.auth.token == "execution-token" From 2fb2423bebaf005506de377340de4f927df04af3 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Sat, 25 Jul 2026 00:32:45 +0000 Subject: [PATCH 4/4] Match surrounding code style --- .../airflow/api_fastapi/execution_api/app.py | 4 +- .../execution_api/routes/callbacks.py | 16 +------ .../api_fastapi/execution_api/security.py | 43 ++++++++++--------- .../versions/head/test_callbacks.py | 6 +-- task-sdk/src/airflow/sdk/api/client.py | 7 +-- .../sdk/execution_time/callback_supervisor.py | 2 +- task-sdk/tests/task_sdk/api/test_client.py | 2 - 7 files changed, 32 insertions(+), 48 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index 73afc9a4f9af8..53be7e5af6769 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -148,7 +148,9 @@ async def dispatch(self, request: Request, call_next): validator: JWTValidator = await services.aget(JWTValidator) claims = await validator.avalidated_claims(token, {}) - # Workload and callback tokens are long-lived (exchanged once, not refreshed here). + # 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 diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py index 031ef87af75cc..b152e98ce9541 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py @@ -18,16 +18,13 @@ from uuid import UUID -import structlog 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.datamodels.token import TIToken from airflow.api_fastapi.execution_api.deps import DepContainer from airflow.api_fastapi.execution_api.security import ( - CurrentTIToken, ExecutionAPIRoute, issue_execution_token, require_auth, @@ -35,8 +32,6 @@ from airflow.models.callback import Callback from airflow.utils.state import CallbackState -log = structlog.get_logger(logger_name=__name__) - router = VersionedAPIRouter( route_class=ExecutionAPIRoute, dependencies=[ @@ -60,14 +55,8 @@ def run_callback( response: Response, session: SessionDep, services=DepContainer, - token: TIToken = CurrentTIToken, ) -> None: - """ - Exchange a single-use callback token for a short-lived execution token. - - Gated on an atomic ``QUEUED -> RUNNING`` transition, so the exchange - succeeds at most once; the new token is returned via ``Refreshed-API-Token``. - """ + """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( @@ -87,11 +76,10 @@ def run_callback( f"Callback {callback_id} is in state {callback.state}; its token can only be " "exchanged once while QUEUED." ), - "previous_state": str(callback.state) if callback.state is not None else None, + "previous_state": callback.state, }, ) callback.state = CallbackState.RUNNING - log.info("Callback token exchanged; marked running", callback_id=str(callback_id)) issue_execution_token(services, response, sub=str(callback_id)) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/security.py b/airflow-core/src/airflow/api_fastapi/execution_api/security.py index b7bb6d9e153ab..dbc3d955c2242 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -85,13 +85,6 @@ VALID_TOKEN_TYPES: frozenset[str] = frozenset(get_args(TokenScope)) -# Maps each ``*:self`` scope to the path param whose value the JWT ``sub`` must match. -SELF_SCOPE_PATH_PARAMS: dict[str, str] = { - "ti:self": "task_instance_id", - "ct:self": "connection_test_id", - "callback:self": "callback_id", -} - _REQUEST_SCOPE_TOKEN_KEY = "ti_token" @@ -189,14 +182,27 @@ async def require_auth( f"Allowed types: {', '.join(sorted(allowed_token_types))}", ) - for scope, path_param in SELF_SCOPE_PATH_PARAMS.items(): - if scope in security_scopes.scopes: - if str(token.id) != str(request.path_params[path_param]): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"Token subject does not match {path_param}", - ) - break + if "ti:self" in security_scopes.scopes: + ti_self_id = str(request.path_params["task_instance_id"]) + if str(token.id) != ti_self_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Token subject does not match task instance ID", + ) + elif "ct:self" in security_scopes.scopes: + ct_self_id = str(request.path_params["connection_test_id"]) + if str(token.id) != ct_self_id: + raise HTTPException( + 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 @@ -205,12 +211,7 @@ async def 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. - - ``JWTReissueMiddleware`` skips ``workload``/``callback`` tokens, so endpoints that - swap them for an execution token must set the header here themselves. - """ + """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"}) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py index 7c95df4641a19..6646b165104bb 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py @@ -33,11 +33,11 @@ def sync_callback(): - """Empty (sync) callable used for unit tests.""" + """Empty (sync) callable used for unit tests""" + pass def _make_queued_callback(session, dag_id="test_dag"): - """Create and persist an ExecutorCallback in QUEUED state.""" callback = ExecutorCallback( callback_def=SyncCallback(sync_callback, kwargs={}), fetch_method=CallbackFetchMethod.IMPORT_PATH, @@ -109,7 +109,7 @@ def test_run_returns_404_for_nonexistent(self, client): @pytest.fixture def _use_real_jwt_bearer(exec_app): - """Remove the mock require_auth override so the real scope validation runs.""" + """Remove the mock require_auth override so the real JWT validation runs end-to-end.""" exec_app.dependency_overrides.pop(require_auth, None) diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index ec6aebf582afe..233e636b0ae39 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -1098,12 +1098,7 @@ def __init__(self, client: Client): self.client = client def run(self, callback_id: uuid.UUID) -> None: - """ - Exchange the single-use callback token for an execution token. - - Must be called before any context read; a second call raises - ``ServerResponseError`` (409). - """ + """Exchange the single-use callback token for an execution token.""" self.client.patch(f"callbacks/{callback_id}/run") diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py index 0b51a5a31580a..6b7650b9e4c63 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -433,7 +433,7 @@ def supervise_callback( else: logger = structlog.get_logger(logger_name="callback").bind() - # Exchange the single-use callback token before any context read; raises on duplicate delivery. + # Swap the single-use callback token for an execution token before any context read. client.callbacks.run(UUID(id)) try: diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index a1274ec08f6d4..dc70c9e0268da 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -2186,7 +2186,6 @@ def handle_request(request: httpx.Request) -> httpx.Response: class TestCallbackOperations: def test_run_exchanges_token(self): - """run() PATCHes the exchange endpoint and adopts the refreshed execution token.""" callback_id = uuid7() def handle_request(request: httpx.Request) -> httpx.Response: @@ -2204,7 +2203,6 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert client.auth.token == "execution-token" def test_run_raises_on_conflict(self): - """A second exchange (callback no longer QUEUED) surfaces as ServerResponseError.""" callback_id = uuid7() def handle_request(request: httpx.Request) -> httpx.Response: