This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1597221692a Validate trigger events in Anthropic deferrable tasks
(#69379)
1597221692a is described below
commit 1597221692aee3180d579412b213d6b9f789be52
Author: Takayoshi Makabe <[email protected]>
AuthorDate: Tue Jul 7 03:59:36 2026 +0900
Validate trigger events in Anthropic deferrable tasks (#69379)
* Validate trigger events in Anthropic deferrable tasks
* Apply suggestions from code review
Co-authored-by: Kaxil Naik <[email protected]>
---------
Co-authored-by: Kaxil Naik <[email protected]>
---
.../src/airflow/providers/anthropic/exceptions.py | 4 +++
.../airflow/providers/anthropic/hooks/anthropic.py | 24 ++++++++++++++++++
.../airflow/providers/anthropic/operators/agent.py | 5 ++--
.../airflow/providers/anthropic/operators/batch.py | 7 +++++-
.../airflow/providers/anthropic/sensors/batch.py | 8 +++++-
.../tests/unit/anthropic/hooks/test_anthropic.py | 29 ++++++++++++++++++++++
.../tests/unit/anthropic/operators/test_agent.py | 23 ++++++++++++++---
.../tests/unit/anthropic/operators/test_batch.py | 18 +++++++++++++-
.../tests/unit/anthropic/sensors/test_batch.py | 18 +++++++++++++-
9 files changed, 125 insertions(+), 11 deletions(-)
diff --git a/providers/anthropic/src/airflow/providers/anthropic/exceptions.py
b/providers/anthropic/src/airflow/providers/anthropic/exceptions.py
index f4e81b868a1..7250d1e848c 100644
--- a/providers/anthropic/src/airflow/providers/anthropic/exceptions.py
+++ b/providers/anthropic/src/airflow/providers/anthropic/exceptions.py
@@ -29,6 +29,10 @@ class AnthropicBatchTimeout(AnthropicError):
"""Raised when an Anthropic Message Batch does not reach a terminal status
in time."""
+class AnthropicTriggerEventError(AnthropicError):
+ """Raised when a deferred task resumes with a missing or malformed trigger
event."""
+
+
class AnthropicAgentSessionError(AnthropicError):
"""Raised when a Managed Agents session terminates or fails."""
diff --git
a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py
b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py
index 4d13240422c..cdd344f8476 100644
--- a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py
+++ b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py
@@ -38,6 +38,7 @@ from airflow.providers.anthropic.exceptions import (
AnthropicBatchJobError,
AnthropicBatchTimeout,
AnthropicError,
+ AnthropicTriggerEventError,
)
from airflow.providers.common.compat.sdk import AirflowSkipException, BaseHook
@@ -149,6 +150,29 @@ def evaluate_session_state(
return False, None, False
+#: Statuses the provider's triggers emit in their terminal event.
+TRIGGER_EVENT_STATUSES = frozenset({"success", "error", "timeout"})
+
+
+def validate_execute_complete_event(event: dict[str, Any] | None = None) ->
dict[str, Any]:
+ """
+ Validate the event a deferred task resumes with, returning it if
well-formed.
+
+ The event crosses the triggerer/worker boundary through the metadata DB,
so a
+ resuming task can receive ``None`` or a status its handlers do not
recognize
+ (version skew, a custom trigger). Both must fail loudly: the
``execute_complete``
+ handlers raise on ``timeout``/``error`` and treat everything else as
success, so
+ an unrecognized status would otherwise silently succeed.
+ """
+ if event is None:
+ raise AnthropicTriggerEventError("Trigger error: event is None")
+ if event.get("status") not in TRIGGER_EVENT_STATUSES:
+ raise AnthropicTriggerEventError(
+ f"Unexpected trigger event status {event.get('status')!r}:
{event!r}"
+ )
+ return event
+
+
def evaluate_batch_counts(
*,
batch_id: str | None,
diff --git
a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py
b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py
index 7b1f5f86e14..22e2860c64d 100644
--- a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py
+++ b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py
@@ -23,7 +23,7 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any
from airflow.providers.anthropic.exceptions import AnthropicAgentSessionError,
AnthropicAgentSessionTimeout
-from airflow.providers.anthropic.hooks.anthropic import AnthropicHook
+from airflow.providers.anthropic.hooks.anthropic import AnthropicHook,
validate_execute_complete_event
from airflow.providers.anthropic.triggers.agent import
AnthropicAgentSessionTrigger
from airflow.providers.common.compat.sdk import BaseOperator, conf
@@ -191,8 +191,7 @@ class AnthropicAgentSessionOperator(BaseOperator):
return session.id
def execute_complete(self, context: Context, event: Any = None) -> str:
- if not event:
- raise AnthropicAgentSessionError("Trigger resumed without an event
payload.")
+ event = validate_execute_complete_event(event)
# The deferred task is a fresh instance; restore the session id from
the event.
self.session_id = event["session_id"]
status = event["status"]
diff --git
a/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py
b/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py
index 7fd6f8f6a88..5ba8722d556 100644
--- a/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py
+++ b/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py
@@ -23,7 +23,11 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any
from airflow.providers.anthropic.exceptions import AnthropicBatchJobError,
AnthropicBatchTimeout
-from airflow.providers.anthropic.hooks.anthropic import AnthropicHook,
evaluate_batch_counts
+from airflow.providers.anthropic.hooks.anthropic import (
+ AnthropicHook,
+ evaluate_batch_counts,
+ validate_execute_complete_event,
+)
from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger
from airflow.providers.common.compat.sdk import BaseOperator, conf
@@ -149,6 +153,7 @@ class AnthropicBatchOperator(BaseOperator):
The deferred task is a fresh instance, so the batch ID is read from
the event,
not ``self.batch_id``.
"""
+ event = validate_execute_complete_event(event)
self.batch_id = event["batch_id"]
status = event["status"]
if status == "timeout":
diff --git
a/providers/anthropic/src/airflow/providers/anthropic/sensors/batch.py
b/providers/anthropic/src/airflow/providers/anthropic/sensors/batch.py
index a5128b5c1b7..dd10d205683 100644
--- a/providers/anthropic/src/airflow/providers/anthropic/sensors/batch.py
+++ b/providers/anthropic/src/airflow/providers/anthropic/sensors/batch.py
@@ -23,7 +23,12 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any
from airflow.providers.anthropic.exceptions import AnthropicBatchJobError,
AnthropicBatchTimeout
-from airflow.providers.anthropic.hooks.anthropic import AnthropicHook,
BatchStatus, evaluate_batch_counts
+from airflow.providers.anthropic.hooks.anthropic import (
+ AnthropicHook,
+ BatchStatus,
+ evaluate_batch_counts,
+ validate_execute_complete_event,
+)
from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger
from airflow.providers.common.compat.sdk import BaseSensorOperator, conf
@@ -107,6 +112,7 @@ class AnthropicBatchSensor(BaseSensorOperator):
super().execute(context)
def execute_complete(self, context: Context, event: Any = None) -> None:
+ event = validate_execute_complete_event(event)
status = event["status"]
if status == "timeout":
raise AnthropicBatchTimeout(event["message"])
diff --git a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
index 44a448b2011..05e39d5a80a 100644
--- a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
+++ b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
@@ -25,6 +25,7 @@ from airflow.providers.anthropic.exceptions import (
AnthropicAgentSessionTimeout,
AnthropicBatchTimeout,
AnthropicError,
+ AnthropicTriggerEventError,
)
from airflow.providers.anthropic.hooks.anthropic import (
DEFAULT_MODEL,
@@ -33,6 +34,7 @@ from airflow.providers.anthropic.hooks.anthropic import (
BatchStatus,
SessionStatus,
evaluate_session_state,
+ validate_execute_complete_event,
)
pytest.importorskip("anthropic")
@@ -76,6 +78,33 @@ class TestSessionStatus:
assert SessionStatus.is_terminal(status) is expected
+class TestValidateTriggerEvent:
+ @pytest.mark.parametrize(
+ ("event", "match"),
+ [
+ pytest.param(None, "event is None", id="none"),
+ pytest.param({}, "Unexpected trigger event status None",
id="missing-status"),
+ pytest.param(
+ {"status": "ended", "batch_id": "b"}, "Unexpected trigger
event status", id="unknown-status"
+ ),
+ ],
+ )
+ def test_invalid_event_raises(self, event, match):
+ with pytest.raises(AnthropicTriggerEventError, match=match):
+ validate_execute_complete_event(event)
+
+ @pytest.mark.parametrize(
+ "event",
+ [
+ pytest.param({"status": "success", "batch_id": "b"}, id="success"),
+ pytest.param({"status": "error", "batch_id": "b", "message":
"boom"}, id="error"),
+ pytest.param({"status": "timeout", "batch_id": "b", "message":
"slow"}, id="timeout"),
+ ],
+ )
+ def test_valid_event_is_returned(self, event):
+ assert validate_execute_complete_event(event) is event
+
+
def _session(status, outcome_results=None):
s = mock.MagicMock()
s.status = status
diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py
b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py
index d461f060f34..3c60e3e0e80 100644
--- a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py
+++ b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py
@@ -21,7 +21,11 @@ from unittest import mock
import pytest
from airflow.exceptions import TaskDeferred
-from airflow.providers.anthropic.exceptions import AnthropicAgentSessionError,
AnthropicAgentSessionTimeout
+from airflow.providers.anthropic.exceptions import (
+ AnthropicAgentSessionError,
+ AnthropicAgentSessionTimeout,
+ AnthropicTriggerEventError,
+)
from airflow.providers.anthropic.hooks.anthropic import AnthropicHook
from airflow.providers.anthropic.operators.agent import
AnthropicAgentSessionOperator
from airflow.providers.anthropic.triggers.agent import
AnthropicAgentSessionTrigger
@@ -201,10 +205,21 @@ class TestExecuteComplete:
op.execute_complete({}, {"status": "timeout", "session_id": "s",
"message": "slow"})
hook.archive_session.assert_called_once_with("s")
- def test_none_event_raises(self):
+ @pytest.mark.parametrize(
+ ("event", "match"),
+ [
+ pytest.param(None, "event is None", id="none"),
+ pytest.param(
+ {"status": "rescheduling", "session_id": "s"},
+ "Unexpected trigger event status",
+ id="unknown-status",
+ ),
+ ],
+ )
+ def test_invalid_event_raises(self, event, match):
op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag",
environment_id="env", message="hi")
- with pytest.raises(AnthropicAgentSessionError, match="without an
event"):
- op.execute_complete({}, None)
+ with pytest.raises(AnthropicTriggerEventError, match=match):
+ op.execute_complete({}, event)
class TestOnKill:
diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
b/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
index b79d1e5b037..656658b3362 100644
--- a/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
+++ b/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
@@ -21,7 +21,11 @@ from unittest import mock
import pytest
from airflow.exceptions import TaskDeferred
-from airflow.providers.anthropic.exceptions import AnthropicBatchJobError,
AnthropicBatchTimeout
+from airflow.providers.anthropic.exceptions import (
+ AnthropicBatchJobError,
+ AnthropicBatchTimeout,
+ AnthropicTriggerEventError,
+)
from airflow.providers.anthropic.hooks.anthropic import AnthropicHook
from airflow.providers.anthropic.operators.batch import AnthropicBatchOperator
from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger
@@ -188,6 +192,18 @@ class TestExecuteComplete:
with pytest.raises(AnthropicBatchJobError, match="failed request"):
op.execute_complete(_context(), event)
+ @pytest.mark.parametrize(
+ "event",
+ [
+ pytest.param(None, id="none"),
+ pytest.param({"status": "ended", "batch_id": "b"},
id="unknown-status"),
+ ],
+ )
+ def test_invalid_event_raises_instead_of_succeeding(self, event):
+ op = AnthropicBatchOperator(task_id="t", requests=REQUESTS)
+ with pytest.raises(AnthropicTriggerEventError):
+ op.execute_complete(_context(), event)
+
class TestOnKill:
@mock.patch.object(AnthropicBatchOperator, "hook",
new_callable=mock.PropertyMock)
diff --git a/providers/anthropic/tests/unit/anthropic/sensors/test_batch.py
b/providers/anthropic/tests/unit/anthropic/sensors/test_batch.py
index 02d071feb06..af5d3ed4657 100644
--- a/providers/anthropic/tests/unit/anthropic/sensors/test_batch.py
+++ b/providers/anthropic/tests/unit/anthropic/sensors/test_batch.py
@@ -21,7 +21,11 @@ from unittest import mock
import pytest
from airflow.exceptions import TaskDeferred
-from airflow.providers.anthropic.exceptions import AnthropicBatchJobError,
AnthropicBatchTimeout
+from airflow.providers.anthropic.exceptions import (
+ AnthropicBatchJobError,
+ AnthropicBatchTimeout,
+ AnthropicTriggerEventError,
+)
from airflow.providers.anthropic.hooks.anthropic import AnthropicHook
from airflow.providers.anthropic.sensors.batch import AnthropicBatchSensor
from airflow.providers.anthropic.triggers.batch import AnthropicBatchTrigger
@@ -102,3 +106,15 @@ class TestAnthropicBatchSensorDeferrable:
sensor = AnthropicBatchSensor(task_id="s", batch_id="b1")
event = {"status": "success", "batch_id": "b1", "request_counts":
{"succeeded": 2}}
assert sensor.execute_complete({}, event) is None
+
+ @pytest.mark.parametrize(
+ "event",
+ [
+ pytest.param(None, id="none"),
+ pytest.param({"status": "ended", "batch_id": "b1"},
id="unknown-status"),
+ ],
+ )
+ def test_execute_complete_invalid_event_raises_instead_of_succeeding(self,
event):
+ sensor = AnthropicBatchSensor(task_id="s", batch_id="b1")
+ with pytest.raises(AnthropicTriggerEventError):
+ sensor.execute_complete({}, event)
