address review feeback

Author: anishgirianish

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

@@ -418,6 +418,10 @@ class JWTGenerator:
 
     kid: str = attrs.field(default=attrs.Factory(_generate_kid, takes_self=True))
     valid_for: float
+    workload_valid_for: float = attrs.field(
+        factory=_conf_factory("execution_api", "jwt_workload_token_expiration_time", fallback="86400"),
+        converter=float,
+    )
     audience: str
     issuer: str | list[str] | None = attrs.field(
         factory=_conf_list_factory("api_auth", "jwt_issuer", first_only=True, fallback=None)
@@ -447,18 +451,6 @@ def signing_arg(self) -> AllowedPrivateKeys | str:
             assert self._secret_key
         return self._secret_key
 
-    def generate_workload_token(self, sub: str) -> str:
-        """Generate a long-lived workload token for executor queues."""
-        from airflow.configuration import conf
-
-        workload_valid_for = conf.getint(
-            "execution_api", "jwt_workload_token_expiration_time", fallback=86400
-        )
-        return self.generate(
-            extras={"sub": sub, "scope": "workload"},
-            valid_for=workload_valid_for,
-        )
-
     def generate(
         self,
         extras: dict[str, Any] | None = None,

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

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import json
+import secrets
 import time
 from contextlib import AsyncExitStack
 from functools import cached_property
@@ -76,6 +77,7 @@ def _jwt_generator():
 
     generator = JWTGenerator(
         valid_for=conf.getint("execution_api", "jwt_expiration_time"),
+        workload_valid_for=conf.getint("execution_api", "jwt_workload_token_expiration_time", fallback=86400),
         audience=conf.get_mandatory_list_value("execution_api", "jwt_audience")[0],
         issuer=conf.get("api_auth", "jwt_issuer", fallback=None),
         # Since this one is used across components/server, there is no point trying to generate one, error
@@ -142,6 +144,11 @@ 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 claims.get("scope") == "workload":
+                        return response
+
                     now = int(time.time())
                     validity = conf.getint("execution_api", "jwt_expiration_time")
                     refresh_when_less_than = max(int(validity * 0.20), 30)
@@ -341,8 +348,14 @@ async def always_allow(request: Request):
             # Cadwyn's versioned sub-apps don't inherit the main app's state,
             # so lookups raise ServiceNotFoundError. This registry provides
             # services needed by routes called during dag.test().
+            #
+            stub_generator = JWTGenerator(
+                secret_key=secrets.token_urlsafe(32),
+                audience="in-process",
+                valid_for=3600,
+            )
             registry = svcs.Registry()
-            registry.register_factory(JWTGenerator, _jwt_generator)
+            registry.register_value(JWTGenerator, stub_generator)
 
             async def _in_process_container(request: Request):
                 async with svcs.Container(registry) as cont:

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

@@ -102,7 +102,7 @@
     },
     response_model_exclude_unset=True,
 )
-async def ti_run(
+def ti_run(
     task_instance_id: UUID,
     ti_run_payload: Annotated[TIEnterRunningPayload, Body()],
     response: Response,
@@ -286,18 +286,18 @@ async def ti_run(
         if ti.next_method:
             context.next_method = ti.next_method
             context.next_kwargs = ti.next_kwargs
-
-        generator: JWTGenerator = await services.aget(JWTGenerator)
-        execution_token = generator.generate(extras={"sub": str(task_instance_id)})
-        response.headers["X-Execution-Token"] = execution_token
-
-        return context
     except SQLAlchemyError:
         log.exception("Error marking Task Instance state as running")
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Database error occurred"
         )
 
+    generator: JWTGenerator = services.get(JWTGenerator)
+    execution_token = generator.generate(extras={"sub": str(task_instance_id)})
+    response.headers["X-Execution-Token"] = execution_token
+
+    return context
+
 
 @ti_id_router.patch(
     "/{task_instance_id}/state",

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

@@ -2074,7 +2074,7 @@ execution_api:
         Seconds until workload JWT tokens expire. These long-lived tokens are sent
         with task workloads to executors and can only call the /run endpoint.
         Set long enough to cover maximum expected queue wait time.
-      version_added: ~
+      version_added: 3.2.0
       type: integer
       example: ~
       default: "86400"

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

@@ -74,7 +74,12 @@ class BaseWorkloadSchema(BaseModel):
 
     @staticmethod
     def generate_token(sub_id: str, generator: JWTGenerator | None = None) -> str:
-        return generator.generate_workload_token(sub_id) if generator else ""
+        if not generator:
+            return ""
+        return generator.generate(
+            extras={"sub": sub_id, "scope": "workload"},
+            valid_for=generator.workload_valid_for,
+        )
 
 
 class BaseDagBundleWorkload(BaseWorkloadSchema, ABC):

airflow-core/tests/unit/api_fastapi/auth/test_tokens.py

@@ -160,48 +160,33 @@ def test_secret_key_with_configured_kid():
         assert header["kid"] == "my-custom-kid"
 
 
-def test_generate_workload_token():
-    """generate_workload_token() produces a token with scope 'workload' and 24h expiry."""
+def test_generate_with_custom_valid_for():
+    """generate() accepts a valid_for override."""
     generator = JWTGenerator(secret_key="test-secret", audience="test", valid_for=60)
+    token = generator.generate(extras={"sub": "user"}, valid_for=3600)
+    claims = jwt.decode(token, "test-secret", algorithms=["HS512"], audience="test")
+    assert claims["exp"] - claims["iat"] == 3600
 
-    with patch.dict(
-        "os.environ",
-        {"AIRFLOW__EXECUTION_API__JWT_WORKLOAD_TOKEN_EXPIRATION_TIME": "86400"},
-    ):
-        token = generator.generate_workload_token(sub="ti-123")
 
+def test_generate_workload_scope_via_extras():
+    """generate() with scope='workload' in extras produces a workload-scoped token."""
+    generator = JWTGenerator(secret_key="test-secret", audience="test", valid_for=60)
+
+    token = generator.generate(extras={"sub": "ti-123", "scope": "workload"}, valid_for=86400)
     claims = jwt.decode(token, "test-secret", algorithms=["HS512"], audience="test")
     assert claims["sub"] == "ti-123"
     assert claims["scope"] == "workload"
-    # Workload token should have ~24h validity, not the generator's default 60s
     assert claims["exp"] - claims["iat"] == 86400
 
 
-def test_generate_with_custom_valid_for():
-    """generate() accepts a valid_for override."""
-    generator = JWTGenerator(secret_key="test-secret", audience="test", valid_for=60)
-    token = generator.generate(extras={"sub": "user"}, valid_for=3600)
-    claims = jwt.decode(token, "test-secret", algorithms=["HS512"], audience="test")
-    assert claims["exp"] - claims["iat"] == 3600
-
-
-def test_workload_token_vs_regular_token_scope():
-    """Regular tokens have no scope, workload tokens have scope 'workload'."""
+def test_regular_token_has_no_scope():
+    """Regular tokens without scope in extras have no scope claim."""
     generator = JWTGenerator(secret_key="test-secret", audience="test", valid_for=60)
 
     regular = generator.generate(extras={"sub": "user"})
     regular_claims = jwt.decode(regular, "test-secret", algorithms=["HS512"], audience="test")
     assert "scope" not in regular_claims
 
-    with patch.dict(
-        "os.environ",
-        {"AIRFLOW__EXECUTION_API__JWT_WORKLOAD_TOKEN_EXPIRATION_TIME": "86400"},
-    ):
-        workload = generator.generate_workload_token(sub="ti-123")
-
-    workload_claims = jwt.decode(workload, "test-secret", algorithms=["HS512"], audience="test")
-    assert workload_claims["scope"] == "workload"
-
 
 @pytest.fixture
 def jwt_generator(ed25519_private_key: Ed25519PrivateKey):

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

@@ -59,9 +59,10 @@ async def mock_jwt_bearer(request: Request):
     with TestClient(app, headers={"Authorization": "Bearer fake"}) as client:
         mock_generator = MagicMock(spec=JWTGenerator)
         mock_generator.generate.return_value = "mock-execution-token"
-        mock_generator.generate_workload_token.return_value = "mock-workload-token"
         lifespan.registry.register_value(JWTGenerator, mock_generator)
 
         yield client
 
+        lifespan.registry.close()
+
     exec_app.dependency_overrides.pop(_jwt_bearer, None)

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -3256,6 +3256,44 @@ def test_invalid_scope_value_rejected(self, client, session, create_task_instanc
         assert resp.status_code == 403
         assert "Invalid token scope" in resp.json()["detail"]
 
+    def test_workload_scope_accepted_on_run_endpoint(
+        self, client, session, create_task_instance, time_machine
+    ):
+        """workload scoped tokens should be accepted on the /run endpoint."""
+        instant = timezone.parse("2024-10-31T12:00:00Z")
+        time_machine.move_to(instant, tick=False)
+
+        ti = create_task_instance(
+            task_id="test_workload_run",
+            state=State.QUEUED,
+            dagrun_state=DagRunState.RUNNING,
+            session=session,
+            start_date=instant,
+            dag_id=str(uuid4()),
+        )
+        session.commit()
+
+        validator = mock.AsyncMock(spec=JWTValidator)
+        validator.avalidated_claims.side_effect = lambda cred, validators: {
+            "sub": str(ti.id),
+            "scope": "workload",
+            "exp": 9999999999,
+            "iat": 1000000000,
+        }
+        lifespan.registry.register_value(JWTValidator, validator)
+
+        resp = client.patch(
+            f"/execution/task-instances/{ti.id}/run",
+            json={
+                "state": "running",
+                "hostname": "test-host",
+                "unixname": "test-user",
+                "pid": 100,
+                "start_date": "2024-10-31T12:00:00Z",
+            },
+        )
+        assert resp.status_code == 200
+
     def test_no_scope_defaults_to_execution(self, client, session, create_task_instance):
         """Tokens without scope claim should default to 'execution'."""
         ti = create_task_instance(task_id="test_no_scope", state=State.RUNNING)

airflow-core/tests/unit/executors/test_workloads.py

@@ -20,9 +20,12 @@
 from pathlib import PurePosixPath
 from uuid import uuid4
 
+import jwt
+
+from airflow.api_fastapi.auth.tokens import JWTGenerator
 from airflow.executors import workloads
 from airflow.executors.workloads import TaskInstance, TaskInstanceDTO
-from airflow.executors.workloads.base import BundleInfo
+from airflow.executors.workloads.base import BaseWorkloadSchema, BundleInfo
 from airflow.executors.workloads.task import ExecuteTask
 
 
@@ -61,3 +64,21 @@ def test_token_excluded_from_workload_repr():
     assert fake_token not in workload_repr, f"JWT token leaked into repr! Found token in: {workload_repr}"
     # But token should still be accessible as an attribute
     assert workload.token == fake_token
+
+
+def test_generate_token_produces_workload_scope():
+    """generate_token should create a JWT with scope 'workload' and workload_valid_for expiry."""
+    generator = JWTGenerator(
+        secret_key="test-secret", audience="test", valid_for=60, workload_valid_for=86400
+    )
+    token = BaseWorkloadSchema.generate_token("ti-123", generator)
+
+    claims = jwt.decode(token, "test-secret", algorithms=["HS512"], audience="test")
+    assert claims["sub"] == "ti-123"
+    assert claims["scope"] == "workload"
+    assert claims["exp"] - claims["iat"] == 86400
+
+
+def test_generate_token_without_generator():
+    """generate_token should return empty string when no generator is provided."""
+    assert BaseWorkloadSchema.generate_token("ti-123", None) == ""

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -302,7 +302,6 @@ def set_instance_attrs(self) -> Generator:
     def mock_executors(self):
         mock_jwt_generator = MagicMock(spec=JWTGenerator)
         mock_jwt_generator.generate.return_value = "mock-token"
-        mock_jwt_generator.generate_workload_token.return_value = "mock-workload-token"
 
         default_executor = mock.MagicMock(name="DefaultExecutor", slots_available=8, slots_occupied=0)
         default_executor.name = ExecutorName(alias="default_exec", module_path="default.exec.module.path")