This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 694e5b50352 Redact secrets from LLMRetryPolicy classification prompts
(#70229)
694e5b50352 is described below
commit 694e5b5035276f5893cf5420cd5fe5414b728ff4
Author: Wei Lee <[email protected]>
AuthorDate: Fri Jul 31 13:35:30 2026 +0800
Redact secrets from LLMRetryPolicy classification prompts (#70229)
---
providers/common/ai/docs/retry_policies.rst | 22 +++++++++--
.../airflow/providers/common/ai/policies/retry.py | 30 ++++++++++++++-
.../ai/tests/unit/common/ai/policies/test_retry.py | 44 ++++++++++++++++++++++
3 files changed, 92 insertions(+), 4 deletions(-)
diff --git a/providers/common/ai/docs/retry_policies.rst
b/providers/common/ai/docs/retry_policies.rst
index a24e5df4481..b498050b685 100644
--- a/providers/common/ai/docs/retry_policies.rst
+++ b/providers/common/ai/docs/retry_policies.rst
@@ -70,7 +70,9 @@ How it works
When a task fails, ``LLMRetryPolicy``:
-1. Sends the exception message to the configured LLM
+1. Sends the exception message to the configured LLM. By default, the message
+ is first masked through Airflow's secrets masker (see ``redact_exception``
+ below) before it is added to the prompt.
2. The LLM classifies the error into a category (``rate_limit``, ``auth``,
``network``, ``data``, ``transient``, ``permanent``)
3. Based on the classification, returns RETRY (with a suggested delay) or FAIL
@@ -156,12 +158,26 @@ Parameters
* - ``timeout``
- 30.0
- Max seconds to wait for the LLM response before falling back.
+ * - ``redact_exception``
+ - True
+ - When ``True``, the exception's string representation is passed through
+ Airflow's secrets masker before being added to the classification
+ prompt. This only masks values already registered via
+ ``mask_secret()`` (e.g. connection passwords Airflow captured while
+ resolving the failing task's connections) -- it is not general-purpose
+ PII detection and will not catch arbitrary sensitive strings that were
+ never registered as secrets. Set to ``False`` only if you are certain
+ your exception messages contain no sensitive data and you need the
+ raw text for accurate classification.
Local LLM support
-----------------
-For environments where exception data must not leave the infrastructure, point
-to a local model via Ollama or vLLM -- see :ref:`howto/self_hosted_models` for
+By default, ``redact_exception`` already masks known secrets before the
+exception data reaches the LLM provider. For environments where exception
+data must not leave your own infrastructure at all -- even in masked form --
+point to a local model via Ollama or vLLM instead, so the classification
+never crosses the network boundary. See :ref:`howto/self_hosted_models` for
general self-hosted connection setup:
.. code-block:: python
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
index f92e4e0d64f..768e16270e2 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
@@ -28,6 +28,8 @@ from typing import TYPE_CHECKING
from pydantic import BaseModel
+from airflow.providers.common.compat.sdk import redact
+
try:
from airflow.sdk.definitions.retry_policy import (
ExceptionRetryPolicy,
@@ -100,6 +102,28 @@ class LLMRetryPolicy(RetryPolicy):
falling back. Defaults to 30s. The LLM provider's own timeout
(e.g. 600s for Anthropic) is much longer; this keeps the retry
decision path fast even when the provider is degraded.
+ :param redact_exception: When ``True`` (the default), the exception's
+ string representation is passed through Airflow's secrets masker
+ (:func:`~airflow.sdk.log.redact`) before being added to the prompt.
+ Set to ``False`` only if you are certain your exception messages
+ contain no sensitive data and you need the raw text for accurate
+ classification.
+
+ .. warning::
+ The exception's string representation is sent to the configured
+ external LLM provider (OpenAI, Anthropic, Bedrock, Vertex, Ollama,
+ etc.) as part of the classification prompt, so it may leak whatever
+ the failing task put in the exception message — connection strings,
+ credential fragments, PII, or other secrets. By default
+ ``_classify()`` runs the message through Airflow's secrets masker
+ (:func:`~airflow.sdk.log.redact`, controlled by ``redact_exception``),
+ which masks values already registered via ``mask_secret()`` (for
+ example, connection passwords Airflow captured while resolving the
+ failing task's connections). This does **not** perform
+ general-purpose PII detection and will not catch arbitrary sensitive
+ strings that were never registered as secrets. You are still
+ responsible for confirming that your task's exception messages are
+ safe to send to a third-party LLM provider.
"""
def __init__(
@@ -109,12 +133,15 @@ class LLMRetryPolicy(RetryPolicy):
instructions: str | None = None,
fallback_rules: list[RetryRule] | None = None,
timeout: float = 30.0,
+ *,
+ redact_exception: bool = True,
) -> None:
self.llm_conn_id = llm_conn_id
self.model_id = model_id
self.instructions = instructions or DEFAULT_INSTRUCTIONS
self.fallback_rules = fallback_rules
self.timeout = timeout
+ self.redact_exception = redact_exception
def evaluate(
self,
@@ -147,10 +174,11 @@ class LLMRetryPolicy(RetryPolicy):
instructions=self.instructions,
)
+ exception_message = redact(str(exception)) if self.redact_exception
else str(exception)
prompt = (
f"Classify this error from a data pipeline task "
f"(attempt {try_number} of {max_tries}):\n\n"
- f"{type(exception).__name__}: {exception}"
+ f"{type(exception).__name__}: {exception_message}"
)
from pydantic_ai.settings import ModelSettings
diff --git a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
index 6f9d976d6f1..5223e883ff7 100644
--- a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
+++ b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
@@ -30,6 +30,7 @@ from airflow.providers.common.ai.policies.retry import (
LLMRetryPolicy,
)
from airflow.sdk.definitions.retry_policy import RetryAction, RetryRule
+from airflow.sdk.log import mask_secret
def _make_mock_agent(category, should_retry, delay=0, reasoning="test"):
@@ -108,6 +109,49 @@ class TestLLMClassifyDecisions:
assert "ValueError: bad column type" in prompt
assert "attempt 2 of 5" in prompt
+ @pytest.mark.enable_redact
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ def test_prompt_redacts_known_secrets(self, mock_hook_cls):
+ secret_value = "super-secret-conn-password"
+ mask_secret(secret_value)
+
+ mock_agent = _make_mock_agent("auth", should_retry=False)
+ mock_hook_cls.return_value.create_agent.return_value = mock_agent
+
+ policy = LLMRetryPolicy(llm_conn_id="test")
+ policy.evaluate(
+ ConnectionError(f"could not authenticate with password
{secret_value}"),
+ try_number=1,
+ max_tries=3,
+ )
+
+ prompt = mock_agent.run_sync.call_args[0][0]
+ assert secret_value not in prompt
+ assert prompt == (
+ "Classify this error from a data pipeline task (attempt 1 of
3):\n\n"
+ "ConnectionError: could not authenticate with password ***"
+ )
+
+ @pytest.mark.enable_redact
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ def test_prompt_keeps_raw_message_when_redaction_disabled(self,
mock_hook_cls):
+
+ secret_value = "super-secret-conn-password"
+ mask_secret(secret_value)
+
+ mock_agent = _make_mock_agent("auth", should_retry=False)
+ mock_hook_cls.return_value.create_agent.return_value = mock_agent
+
+ policy = LLMRetryPolicy(llm_conn_id="test", redact_exception=False)
+ policy.evaluate(
+ ConnectionError(f"could not authenticate with password
{secret_value}"),
+ try_number=1,
+ max_tries=3,
+ )
+
+ prompt = mock_agent.run_sync.call_args[0][0]
+ assert secret_value in prompt
+
@patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
def test_custom_instructions_forwarded_to_agent(self, mock_hook_cls):
mock_hook_cls.return_value.create_agent.return_value =
_make_mock_agent("x", False)
