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 018169c5489 Reject durable replay and HITL review when an agent holds
a SandboxToolset (#73529)
018169c5489 is described below
commit 018169c5489ab1d5ea51391248feb2ce9d289cef
Author: Kaxil Naik <[email protected]>
AuthorDate: Tue Sep 22 11:35:07 2026 +0100
Reject durable replay and HITL review when an agent holds a SandboxToolset
(#73529)
A sandbox is destroyed when the agent run ends, so two operator features
that
assume otherwise produced wrong answers instead of errors. durable=True
replayed
cached tool results such as 'wrote the file' against a sandbox that no
longer
existed, then ran the first cache miss against a fresh empty one.
enable_hitl_review=True regenerated after feedback in a second run whose
message history described files the first run's sandbox had held.
AgentOperator
now refuses both at construction, looking inside prefixed, filtered and
combined
toolsets and inside Toolset capabilities, and names the two ways out.
A backend raising a recoverable SandboxError from create() used to escape
the
toolset's error mapping and fail the task by accident. It now fails it on
purpose: the model has no input into provisioning, so nothing it retries can
help, and Airflow's own retry is the right one. The create() contract says
so.
Also covers the liveness probe branch where a non-Modal error keeps the
handle.
---
providers/common/ai/docs/sandbox/configuration.rst | 27 ++++--
providers/common/ai/docs/sandbox/index.rst | 4 +-
.../airflow/providers/common/ai/operators/agent.py | 53 ++++++++++-
.../airflow/providers/common/ai/sandbox/base.py | 10 +-
.../providers/common/ai/toolsets/sandbox.py | 10 ++
.../airflow/providers/common/ai/utils/toolsets.py | 64 +++++++++++++
.../tests/unit/common/ai/operators/test_agent.py | 105 +++++++++++++++++++++
.../ai/tests/unit/common/ai/sandbox/test_modal.py | 42 +++++++++
.../tests/unit/common/ai/toolsets/test_sandbox.py | 55 ++++++++++-
.../ai/tests/unit/common/ai/utils/test_toolsets.py | 66 +++++++++++++
10 files changed, 423 insertions(+), 13 deletions(-)
diff --git a/providers/common/ai/docs/sandbox/configuration.rst
b/providers/common/ai/docs/sandbox/configuration.rst
index 379fa064f45..fa375be32aa 100644
--- a/providers/common/ai/docs/sandbox/configuration.rst
+++ b/providers/common/ai/docs/sandbox/configuration.rst
@@ -239,20 +239,33 @@ lifetime. Human output review is not part of the run: it
starts after the agent
has finished and the sandbox has been destroyed, so a review pause costs no
sandbox time and keeps no files.
+**A failed provisioning fails the task.** The model has no input into
``create``:
+it takes only the spec, which is fixed in the Dag file. So whatever a backend
+raises while provisioning is treated as terminal, even a ``SandboxError`` that
+would have been a model retry from a tool call. The task fails and Airflow's
own
+retry attempts the provisioning again, instead of the model spending its retry
+budget on an image tag or a credential it cannot see.
+
**Nothing survives the run.** A task retry starts from an empty sandbox, and so
-does every other attempt. Two operator features assume otherwise and must not
be
-combined with a sandbox today, because neither is rejected:
+does every other attempt. Two operator features assume otherwise, and
+``AgentOperator`` refuses each of them at construction when any toolset,
including
+one nested inside ``.prefixed()``, a combined toolset or a ``Toolset``
capability,
+is a ``SandboxToolset``:
- ``durable=True`` caches each tool result and replays it on a retry without
- calling the backend, so a replayed ``write_file`` reports success while no
- sandbox exists, and the first call that misses the cache runs against a fresh
- empty one. The model is handed a filesystem that does not match what it was
just
- told, and nothing raises.
+ calling the backend, so a replayed ``write_file`` would report success while
no
+ sandbox exists, and the first call that misses the cache would run against a
+ fresh empty one.
- ``enable_hitl_review=True`` regenerates after reviewer feedback by starting a
second agent run, and the first run's sandbox was destroyed when that run
ended.
- The regenerated agent gets an empty sandbox while its own history describes
+ The regenerated agent would get an empty sandbox while its own history
describes
files it wrote earlier.
+The two ways out are dropping the flag, or moving the sandbox work into its own
+task and keeping the durable or reviewed agent free of sandbox tools. A toolset
+resolved per run from a callable cannot be inspected when the operator is
built,
+so it is the one composition the check does not see.
+
Cost and operations
-------------------
diff --git a/providers/common/ai/docs/sandbox/index.rst
b/providers/common/ai/docs/sandbox/index.rst
index 32908cfc61e..daa15497efd 100644
--- a/providers/common/ai/docs/sandbox/index.rst
+++ b/providers/common/ai/docs/sandbox/index.rst
@@ -385,8 +385,8 @@ is the list to read before designing a Dag around an agent
with a sandbox.
- **A credential cannot come from a connection or a secrets backend**; the
spec is
fixed at parse time, and anything injected is readable by the model.
:ref:`Credentials <sandbox-credentials>`.
-- **Do not combine with** ``durable=True`` **or** ``enable_hitl_review=True``.
- Neither is rejected today. :ref:`Lifecycle <sandbox-lifecycle>`.
+- **Cannot be combined with** ``durable=True`` **or**
``enable_hitl_review=True``;
+ ``AgentOperator`` raises. :ref:`Lifecycle <sandbox-lifecycle>`.
- **A run that outlives** ``sandbox_timeout`` **fails the task.**
:ref:`Lifecycle <sandbox-lifecycle>`.
- **The hostname allowlist is a weak control** and refused unless opted into.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
index 6d3b1f71fa3..4038edf9f12 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
@@ -26,6 +26,7 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any, ClassVar
from pydantic import BaseModel
+from pydantic_ai.capabilities import Toolset
from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin
@@ -33,8 +34,10 @@ from airflow.providers.common.ai.observability import (
build_run_identity_attributes,
stamp_identity_on_agent_spans,
)
+from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset
from airflow.providers.common.ai.utils.logging import log_run_summary,
wrap_toolsets_for_logging
from airflow.providers.common.ai.utils.output_type import
rehydrate_pydantic_output
+from airflow.providers.common.ai.utils.toolsets import find_toolset
from airflow.providers.common.ai.utils.usage import coerce_usage_limits
from airflow.providers.common.compat.sdk import (
AirflowOptionalProviderFeatureException,
@@ -199,6 +202,9 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
retry; put tools you need replayed in ``toolsets=``. Provider-native
capabilities such as ``WebSearch`` and ``Thinking`` execute inside the
model call and are covered by model-response caching.
+ Cannot be combined with a ``SandboxToolset`` (raises): a sandbox is
+ destroyed when the run ends, so replayed tool results would describe
+ files that no longer exist.
:param code_mode: When ``True``, wraps the agent's tools in a single
``run_code`` tool powered by the Monty sandbox (pydantic-ai-harness
``CodeMode``). Instead of one model round-trip per tool call, the model
@@ -234,7 +240,10 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
iterative review loop after the first generation. A human reviewer
can approve, reject, or request changes via the plugin's REST API
at ``/hitl-review`` or through the **HITL Review** extra link
- on the task instance. Default ``False``.
+ on the task instance. Default ``False``. Cannot be combined with a
+ ``SandboxToolset`` (raises): regeneration after feedback is a second
+ run, which starts from an empty sandbox while its history describes
+ the first run's files.
:param max_hitl_iterations: Maximum outputs shown to the reviewer (1 =
initial output). When the reviewer requests changes at
iteration >= this limit, the task fails with ``HITLMaxIterationsError``
@@ -350,11 +359,52 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
# surface the final message history.
raise ValueError("message_history and enable_hitl_review=True
cannot be used together.")
+ if durable or enable_hitl_review:
+ self._reject_sandbox_without_continuity(durable=durable,
enable_hitl_review=enable_hitl_review)
+
self.enable_hitl_review = enable_hitl_review
self.max_hitl_iterations = max_hitl_iterations
self.hitl_timeout = hitl_timeout
self.hitl_poll_interval = hitl_poll_interval
+ def _reject_sandbox_without_continuity(self, *, durable: bool,
enable_hitl_review: bool) -> None:
+ """
+ Refuse a ``SandboxToolset`` under a feature that assumes the sandbox
outlives the run.
+
+ A sandbox is provisioned on the first tool call and destroyed when the
run
+ ends, so nothing in it survives into a retry or a second run. Two
features
+ assume otherwise, and each produces a wrong answer rather than an
error:
+
+ * ``durable=True`` replays cached tool results on a retry without
calling the
+ backend, so a replayed ``write_file`` reports success while no
sandbox exists,
+ and the first call that misses the cache runs against a fresh, empty
one.
+ * ``enable_hitl_review=True`` regenerates after reviewer feedback by
starting a
+ second agent run, which gets an empty sandbox while its message
history still
+ describes the files the first run wrote.
+
+ The toolset is looked for inside wrappers and combinations
(``.prefixed()``,
+ ``.filtered()``, several toolsets passed together) and inside
``Toolset``
+ capabilities, since those are the compositions the documentation
recommends.
+ A toolset resolved per run from a callable cannot be inspected here.
+ """
+ candidates = list(self.toolsets or [])
+ for capability in self.agent_params.get("capabilities") or ():
+ if isinstance(capability, Toolset) and not
callable(capability.toolset):
+ candidates.append(capability.toolset)
+ if find_toolset(candidates, SandboxToolset) is None:
+ return
+ flag = "durable=True" if durable else "enable_hitl_review=True"
+ why = (
+ "cached tool results would be replayed against a sandbox that no
longer exists"
+ if durable
+ else "a regenerated run would start from an empty sandbox while
its history describes "
+ "files from the first run"
+ )
+ raise ValueError(
+ f"{flag} cannot be used with a SandboxToolset: {why}. "
+ f"Drop {flag}, or move the sandbox work into its own task."
+ )
+
@cached_property
def llm_hook(self) -> PydanticAIHook:
"""Return PydanticAIHook for the configured LLM connection."""
@@ -418,7 +468,6 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
# pydantic-ai (and the pydantic-ai-importing CachingToolset) are
imported
# lazily to keep them out of DAG-parse-time imports, matching
# ``_build_durable_toolsets`` and the rest of this module.
- from pydantic_ai.capabilities import Toolset
from pydantic_ai.toolsets.abstract import AbstractToolset
from airflow.providers.common.ai.durable.caching_toolset import
CachingToolset
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py
b/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py
index 34a58ff1778..d43a0cdf3e8 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py
@@ -36,7 +36,9 @@ class SandboxError(Exception):
A sandbox operation failed in a way the agent may be able to work around.
The toolset turns this into a ``ModelRetry`` so the model can adjust and
try
- again within the run (a bad path, a command the image cannot run).
+ again within the run (a bad path, a command the image cannot run). Raised
+ from :meth:`SandboxBackend.create` it is treated as terminal instead, since
+ the model cannot influence provisioning.
"""
@@ -171,6 +173,12 @@ class SandboxBackend(ABC):
backend cannot enforce, rather than provisioning something weaker than
was asked for. It is terminal rather than recoverable because it states
a configuration fact the model cannot see and cannot fix by retrying.
+
+ Every failure raised here is terminal, whichever class carries it. The
+ model has no input into provisioning, so a :class:`SandboxError` from
+ ``create`` is not something it can work around; the toolset re-raises
one
+ as :class:`SandboxTerminalError` and fails the task, so Airflow's retry
+ attempts the provisioning again.
"""
@abstractmethod
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
index 15870bb657c..d06288cd512 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
@@ -280,6 +280,16 @@ class SandboxToolset(AbstractToolset[Any]):
# __aexit__ can destroy a sandbox created during cancellation.
self._sandbox = await create_task
raise
+ except SandboxTerminalError:
+ raise
+ except SandboxError as e:
+ # Provisioning takes only the spec, which the model cannot see or
change,
+ # so no retry the model makes can turn a failed create into a
working
+ # sandbox. Fail the task and let Airflow's retry try the
provisioning
+ # again, rather than spending the model's retry budget on it.
+ raise SandboxTerminalError(
+ f"Could not provision a sandbox on backend
{self._backend.name!r}: {e}"
+ ) from e
else:
self._sandbox = sandbox
return sandbox
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/toolsets.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/toolsets.py
new file mode 100644
index 00000000000..3df255aa816
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/toolsets.py
@@ -0,0 +1,64 @@
+#
+# 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.
+"""Walk a pydantic-ai toolset and the toolsets it wraps or combines."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai.toolsets.combined import CombinedToolset
+from pydantic_ai.toolsets.wrapper import WrapperToolset
+
+if TYPE_CHECKING:
+ from collections.abc import Iterable, Iterator
+
+ from pydantic_ai.toolsets.abstract import AbstractToolset
+
+
+def iter_toolsets(toolset: AbstractToolset[Any]) ->
Iterator[AbstractToolset[Any]]:
+ """
+ Yield ``toolset`` and every toolset nested inside it, outermost first.
+
+ pydantic-ai composes toolsets by wrapping: ``.prefixed()``,
``.filtered()``,
+ ``.prepared()`` and ``.renamed()`` each return a :class:`WrapperToolset`
around
+ the original, and passing several toolsets together produces a
+ :class:`CombinedToolset`. A check that only looks at the top level
therefore
+ misses a toolset the moment an author composes it, which the documentation
+ recommends for running two sandboxes on one agent. This walks both shapes.
+
+ A toolset resolved per run from a callable (``DynamicToolset``) has nothing
+ to walk until the run starts, so it is yielded as itself and its contents
are
+ not inspected.
+ """
+ stack: list[AbstractToolset[Any]] = [toolset]
+ while stack:
+ current = stack.pop()
+ yield current
+ if isinstance(current, WrapperToolset):
+ stack.append(current.wrapped)
+ elif isinstance(current, CombinedToolset):
+ stack.extend(reversed(current.toolsets))
+
+
+def find_toolset(toolsets: Iterable[AbstractToolset[Any]], cls: type) ->
AbstractToolset[Any] | None:
+ """Return the first toolset of type ``cls`` in ``toolsets``, looking
inside wrappers, or ``None``."""
+ for toolset in toolsets:
+ for nested in iter_toolsets(toolset):
+ if isinstance(nested, cls):
+ return nested
+ return None
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index a5d49b11926..89411f55992 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -37,6 +37,7 @@ from pydantic_ai.messages import (
UserPromptPart,
)
from pydantic_ai.models.function import AgentInfo, FunctionModel
+from pydantic_ai.toolsets.combined import CombinedToolset
from pydantic_ai.toolsets.function import FunctionToolset
from pydantic_ai.usage import RequestUsage, UsageLimits
@@ -45,7 +46,9 @@ from airflow.providers.common.ai.durable.caching_toolset
import CachingToolset
from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
from airflow.providers.common.ai.durable.storage import DurableStorage
from airflow.providers.common.ai.operators.agent import AgentOperator,
HITLReviewLink, _build_code_mode
+from airflow.providers.common.ai.sandbox.base import SandboxBackend
from airflow.providers.common.ai.toolsets.logging import LoggingToolset
+from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_3_PLUS
@@ -1212,6 +1215,108 @@ class TestAgentOperatorHITLArgumentChecks:
)
+class _NoopBackend(SandboxBackend):
+ """A backend that is never reached: these tests stop at the operator's
constructor."""
+
+ name = "noop"
+
+ def create(self, *, spec=None):
+ raise AssertionError("constructor guards must not provision")
+
+ def run_command(self, sandbox, command, *, timeout, max_output_bytes):
+ raise AssertionError("constructor guards must not run commands")
+
+ def destroy(self, sandbox):
+ pass
+
+
+def _sandbox_toolset():
+ return SandboxToolset(_NoopBackend())
+
+
+class TestAgentOperatorSandboxContinuityGuards:
+ """
+ A sandbox is destroyed when the run ends, so features that replay or rerun
against
+ it are refused up front rather than producing answers about files that are
gone.
+ """
+
+ @pytest.mark.parametrize(
+ "toolsets",
+ [
+ pytest.param([_sandbox_toolset()], id="direct"),
+ pytest.param([_sandbox_toolset().prefixed("box")], id="prefixed"),
+ pytest.param([_sandbox_toolset().filtered(lambda ctx, tool:
True)], id="filtered"),
+ pytest.param([CombinedToolset([FunctionToolset(),
_sandbox_toolset()])], id="combined"),
+ pytest.param(
+ [CombinedToolset([FunctionToolset(),
_sandbox_toolset().prefixed("a")]).prefixed("b")],
+ id="nested_twice",
+ ),
+ ],
+ )
+ def
test_durable_with_a_sandbox_toolset_is_rejected_wherever_it_hides(self,
toolsets):
+ with pytest.raises(ValueError, match="durable=True cannot be used with
a SandboxToolset"):
+ AgentOperator(task_id="t", prompt="p", llm_conn_id="c",
durable=True, toolsets=toolsets)
+
+ def test_durable_with_a_sandbox_in_a_toolset_capability_is_rejected(self):
+ """Tools reaching the agent through ``capabilities=[Toolset(...)]``
are durably cached
+ too, so the same replay-against-nothing applies to them."""
+ with pytest.raises(ValueError, match="durable=True cannot be used with
a SandboxToolset"):
+ AgentOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ durable=True,
+ agent_params={"capabilities": [Toolset(_sandbox_toolset())]},
+ )
+
+ def
test_a_callable_toolset_capability_cannot_be_inspected_and_passes(self):
+ """A factory resolved per run has no concrete toolset to look inside
at parse time."""
+ op = AgentOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ durable=True,
+ agent_params={"capabilities": [Toolset(lambda ctx:
_sandbox_toolset())]},
+ )
+ assert op.durable is True
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible
with Airflow >= 3.1.0"
+ )
+ def test_hitl_review_with_a_sandbox_toolset_is_rejected(self):
+ with pytest.raises(ValueError, match="enable_hitl_review=True cannot
be used with a SandboxToolset"):
+ AgentOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ enable_hitl_review=True,
+ toolsets=[_sandbox_toolset().prefixed("box")],
+ )
+
+ def test_the_message_names_the_ways_out(self):
+ with pytest.raises(ValueError, match="Drop durable=True, or move the
sandbox work into its own task"):
+ AgentOperator(
+ task_id="t", prompt="p", llm_conn_id="c", durable=True,
toolsets=[_sandbox_toolset()]
+ )
+
+ @pytest.mark.parametrize("flag", [{"durable": True},
{"enable_hitl_review": True}])
+ def test_the_flags_stay_usable_without_a_sandbox(self, flag):
+ if "enable_hitl_review" in flag and not AIRFLOW_V_3_1_PLUS:
+ pytest.skip("Human in the loop is only compatible with Airflow >=
3.1.0")
+ op = AgentOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ toolsets=[FunctionToolset().prefixed("fn"),
CombinedToolset([FunctionToolset()])],
+ **flag,
+ )
+ assert op.toolsets is not None
+
+ def test_a_sandbox_toolset_without_either_flag_is_fine(self):
+ op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c",
toolsets=[_sandbox_toolset()])
+ assert isinstance(op.toolsets[0], SandboxToolset)
+
+
class TestAgentOperatorRunIdentity:
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_run_id_and_usage_pushed_to_xcom(self, mock_hook_cls,
make_mock_run_result):
diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
b/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
index 11cea17ede3..41553d5b9a8 100644
--- a/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
+++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
@@ -19,6 +19,7 @@ from __future__ import annotations
import builtins
import importlib
import json
+import logging
import sys
import threading
@@ -1052,6 +1053,47 @@ class TestMalformedHelperReply:
assert "run_command" not in str(caught.value)
+class TestLivenessProbe:
+ """
+ The probe after a deadline decides whether the handle is kept.
+
+ Only a Modal error is evidence that the sandbox is gone. Anything else
raised by the
+ probe says nothing about the sandbox, and the command it follows has
already produced
+ its output, so the handle is kept and the result is reported as a plain
timeout.
+ """
+
+ def test_a_non_modal_error_in_the_probe_keeps_the_handle(self, backend,
fake, caplog):
+ caplog.set_level(logging.DEBUG,
logger="airflow.providers.common.ai.sandbox.modal")
+ handle, sandbox = _created(backend, fake)
+ sandbox.process = FakeProcess(returncode=-1)
+ # Let the command through, then fail only the probe that follows it.
+ sandbox.exec_errors = [None, RuntimeError("client hiccup")]
+
+ result = backend.run_command(handle, "sleep 5", timeout=2,
max_output_bytes=1024)
+
+ assert result.timed_out is True
+ assert result.sandbox_terminated is False
+ assert "Could not check whether Modal sandbox" in caplog.text
+ # The handle is still cached, so the next command needs no lookup and
no create.
+ fake.Sandbox.from_id_error = AssertionError("a kept handle must not be
looked up again")
+ sandbox.process = FakeProcess(returncode=0)
+ assert backend.run_command(handle, "true", timeout=2,
max_output_bytes=1024).exit_code == 0
+ assert len(fake.Sandbox.by_id) == 1
+
+ def test_a_modal_error_in_the_probe_drops_the_handle(self, backend, fake):
+ handle, sandbox = _created(backend, fake)
+ sandbox.process = FakeProcess(returncode=-1)
+ sandbox.exec_errors = [None, fake.exception.NotFoundError("gone")]
+
+ result = backend.run_command(handle, "sleep 5", timeout=2,
max_output_bytes=1024)
+
+ assert result.sandbox_terminated is True
+ # The cached handle is gone, so a later call has to look the sandbox
up again.
+ fake.Sandbox.from_id_error = fake.exception.NotFoundError("gone")
+ with pytest.raises(SandboxTerminalError):
+ backend.run_command(handle, "true", timeout=2,
max_output_bytes=1024)
+
+
class TestDestroy:
def test_terminates_without_waiting(self, backend, fake):
"""``terminate(wait=True)`` took 31s when measured, and this runs in
every teardown."""
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
index e248047e5c6..c05055c68dd 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
@@ -43,7 +43,14 @@ class _RecordingBackend(SandboxBackend):
name = "rec"
- def __init__(self, *, destroy_error: Exception | None = None,
run_result=None, run_error=None):
+ def __init__(
+ self,
+ *,
+ destroy_error: Exception | None = None,
+ run_result=None,
+ run_error=None,
+ create_error: Exception | None = None,
+ ):
self.created: list[SandboxSpec | None] = []
self.destroyed: list[str] = []
self.commands: list[tuple[str, str, float, int]] = []
@@ -52,11 +59,14 @@ class _RecordingBackend(SandboxBackend):
self.destroy_error = destroy_error
self.run_result = run_result
self.run_error = run_error
+ self.create_error = create_error
self.read_payload = b""
self.read_error: Exception | None = None
def create(self, *, spec: SandboxSpec | None = None) -> str:
self.created.append(spec)
+ if self.create_error is not None:
+ raise self.create_error
return f"box-{len(self.created)}"
def run_command(self, sandbox, command, *, timeout, max_output_bytes):
@@ -445,6 +455,49 @@ class TestErrorMapping:
with pytest.raises(SandboxTerminalError):
await _call(ts, "run_command", {"command": "x"})
+ @pytest.mark.asyncio
+ async def test_a_recoverable_error_from_create_is_terminal(self):
+ # The model has no say in provisioning, so a retry cannot fix a failed
+ # create. It must fail the task rather than reach the model as a
ModelRetry.
+ backend = _RecordingBackend(create_error=SandboxError("image pull
timed out"))
+ ts = SandboxToolset(backend)
+
+ async with ts:
+ with pytest.raises(
+ SandboxTerminalError, match="Could not provision.*image pull
timed out"
+ ) as caught:
+ await _call(ts, "run_command", {"command": "x"})
+
+ assert isinstance(caught.value.__cause__, SandboxError)
+ assert backend.destroyed == [], "nothing was provisioned, so nothing
is destroyed"
+
+ @pytest.mark.asyncio
+ async def test_a_terminal_error_from_create_propagates_unchanged(self):
+ error = SandboxTerminalError("credentials rejected")
+ ts = SandboxToolset(_RecordingBackend(create_error=error))
+
+ async with ts:
+ with pytest.raises(SandboxTerminalError) as caught:
+ await _call(ts, "run_command", {"command": "x"})
+
+ assert caught.value is error
+
+ @pytest.mark.asyncio
+ async def test_a_failed_create_is_retried_by_the_next_call(self):
+ # A create that raised left no sandbox behind, so the next tool call in
+ # the same run must try provisioning again rather than reuse a dead
task.
+ backend = _RecordingBackend(create_error=SandboxError("transient"))
+ ts = SandboxToolset(backend)
+
+ async with ts:
+ with pytest.raises(SandboxTerminalError):
+ await _call(ts, "run_command", {"command": "x"})
+ backend.create_error = None
+ assert await _call(ts, "run_command", {"command": "x"})
+
+ assert len(backend.created) == 2
+ assert backend.destroyed == ["box-2"]
+
@pytest.mark.asyncio
async def test_unknown_tool_raises(self):
ts = SandboxToolset(_RecordingBackend())
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_toolsets.py
b/providers/common/ai/tests/unit/common/ai/utils/test_toolsets.py
new file mode 100644
index 00000000000..e9a5bf2f8c7
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_toolsets.py
@@ -0,0 +1,66 @@
+#
+# 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 pydantic_ai.toolsets.combined import CombinedToolset
+from pydantic_ai.toolsets.function import FunctionToolset
+
+from airflow.providers.common.ai.utils.toolsets import find_toolset,
iter_toolsets
+
+
+class _Marker(FunctionToolset):
+ """A toolset type worth finding."""
+
+
+class TestIterToolsets:
+ def test_a_plain_toolset_yields_itself(self):
+ ts = FunctionToolset()
+ assert list(iter_toolsets(ts)) == [ts]
+
+ def test_wrappers_are_unwrapped_outermost_first(self):
+ inner = FunctionToolset()
+ prefixed = inner.prefixed("a")
+ filtered = prefixed.filtered(lambda ctx, tool: True)
+
+ assert list(iter_toolsets(filtered)) == [filtered, prefixed, inner]
+
+ def test_a_combined_toolset_yields_its_members_in_order(self):
+ a, b = FunctionToolset(), FunctionToolset()
+ combined = CombinedToolset([a, b])
+
+ assert list(iter_toolsets(combined)) == [combined, a, b]
+
+ def
test_wrappers_inside_combinations_inside_wrappers_are_all_reached(self):
+ leaf = _Marker()
+ composed = CombinedToolset([FunctionToolset(),
leaf.prefixed("x")]).prefixed("y")
+
+ assert leaf in list(iter_toolsets(composed))
+
+
+class TestFindToolset:
+ def test_returns_the_nested_instance(self):
+ leaf = _Marker()
+ toolsets = [FunctionToolset(), CombinedToolset([leaf.prefixed("x")])]
+
+ assert find_toolset(toolsets, _Marker) is leaf
+
+ def test_returns_none_when_absent(self):
+ assert find_toolset([FunctionToolset().prefixed("x")], _Marker) is None
+
+ def test_empty_input(self):
+ assert find_toolset([], _Marker) is None
