Claim callbacks through the Execution API before running them - #70927
Claim callbacks through the Execution API before running them#70927potiuk wants to merge 2 commits into
Conversation
A CeleryExecutor worker runs an ExecuteCallback workload by importing the
callback path and calling it in a supervised subprocess. Unlike a task, a
callback made no authenticated call to the Execution API before doing so, so the
workload token minted for it in ExecuteCallback.make was never redeemed anywhere
and the callback executed whether or not the token was valid.
Add POST /execution/callbacks/{callback_id}/run, scoped to token:workload +
cb:self, which atomically transitions the callback QUEUED -> RUNNING, and call it
from supervise_callback before the subprocess starts. The server validates the
token (workload scope, subject == this callback id), so a missing, forged, or
mismatched token is rejected before the callback body runs. The transition is
single-shot, so a redelivered or replayed message that reaches an already-claimed
callback is refused rather than run twice, and it gives callbacks the RUNNING
state the workload docstring already describes.
cb:self mirrors ct:self for connection tests; the supervisor ordering mirrors
supervise_connection_test. The new endpoint is versioned as 2026-08-01 so older
clients continue to negotiate an API version in which it does not exist.
|
@seanghaeli - weren't you working on this a while ago, or was that a different workload token? Found it. I think these two PRs are going to collide, or at least overlap: #71192 |
| }, | ||
| ) | ||
|
|
||
| callback.state = CallbackState.RUNNING |
There was a problem hiding this comment.
Doesn't this conflict with the executors which already change the state from QUEUED to RUNNING?
LocalExecutor queues the RUNNING event before calling run_workload(), so it would be RUNNING before it gets here and this if not QUEUED block would throw an exception, failing the callback by default before it ever gets run. The AWS executors all set the state to RUNNING in self.running_state() as well. I'm pretty sure Celery is the exception here and every other executor would start failing.
I suspect the executors need to be updated to not also change the state, or this if queued check needs to be expanded to if not terminal state, maybe?
|
Yes, I think this PR fully overlaps with #71192. Since Ash has already taken a look at it, let's land that one? In the linked PR, the token is more narrowly scoped to just the A gap in mine though is that it doesn't include API versioning bundle, so I'll pull that over. |
|
Added the versioning bundle code to #71192 |
3 participants
A CeleryExecutor worker runs an
ExecuteCallbackworkload by importing thecallback path and calling it in a supervised subprocess. Unlike a task or a
connection test, a callback made no authenticated call to the Execution API
before doing so — so the workload token minted for it in
ExecuteCallback.makewas never redeemed anywhere, and the callback executed whether or not the token
was valid. This closes that gap and gives callbacks the
RUNNINGstate theworkload state machine already describes (
QUEUED → RUNNING → SUCCESS/FAILED).Change
POST /execution/callbacks/{callback_id}/run, scopedtoken:workload+cb:self, that atomically transitions the callbackQUEUED → RUNNING.supervise_callbackcalls it before starting the subprocess — mirroringsupervise_connection_test, which calls the server before touching credentials.cb:selfenforcement mirrors the existingct:selffor connection tests.2026-08-01, so a client negotiating an older API version behavesas though the endpoint does not exist.
The server validates the token (workload scope, subject == this callback id), so
a missing or mismatched token is refused before the callback body runs. The
transition is single-shot, so a redelivered/replayed message that reaches an
already-claimed callback is refused rather than run twice.
Tests
test_callbacks.py— claim transitions QUEUED→RUNNING; 404; 422; and 409 on acallback already RUNNING/terminal (replay protection)
TestCbSelfScopeEnforcement— a token minted for another callback is rejected (403)test_token_scope_boundaries— the route is registered workload-onlyTestSuperviseCallbackClaimsBeforeExecuting— the claim runs before thesubprocess, and a rejected claim prevents it from starting. Verified: removing
the claim call makes both fail.
@vincbeck @o-nikolas — flagging you as the security-team folks most across this
area. Please pull in @ferruzzi and anyone else who worked on the workload /
callback execution path if I have missed them.
Was generative AI tooling used to co-author this PR?
Generated-by: Claude Opus 5 (1M context) following the guidelines at
https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions