Source code for tests.system.common.ai.example_sandbox_toolset_boat
# 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.
"""End-to-end system test for SandboxToolset with Boat."""
from __future__ import annotations
import os
from datetime import UTC, datetime
from airflow.providers.common.compat.sdk import dag as airflow_dag, task
[docs]
DAG_ID = f"common_ai_sandbox_toolset_boat_{ENV_ID}" if ENV_ID else "common_ai_sandbox_toolset_boat"
@airflow_dag(
dag_id=DAG_ID,
schedule="@once",
start_date=datetime(2024, 1, 1, tzinfo=UTC),
catchup=False,
tags=["common.ai", "sandbox", "boat", "system_test"],
)
[docs]
def example_sandbox_toolset_boat():
@task
def run_sandbox_agent() -> str:
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
from airflow.providers.common.ai.sandbox import BoatSandboxBackend, SandboxSpec
from airflow.providers.common.ai.toolsets import SandboxToolset
class RecordingBackend(BoatSandboxBackend):
"""Remembers the sandboxes it created and destroyed, so teardown can be asserted."""
def __init__(self, **kwargs) -> None:
super().__init__(**kwargs)
self.created: list[str] = []
self.destroyed: list[str] = []
def create(self, *, spec=None) -> str:
sandbox = super().create(spec=spec)
self.created.append(sandbox)
return sandbox
def destroy(self, sandbox: str) -> None:
super().destroy(sandbox)
self.destroyed.append(sandbox)
def model_function(messages: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
returns = [
part.content
for message in messages
for part in message.parts
if part.part_kind == "tool-return"
]
if not returns:
return ModelResponse(
parts=[
ToolCallPart(
tool_name="write_file",
args={"path": STATE_PATH, "content": MARKER},
tool_call_id="write",
)
]
)
if "Wrote" not in str(returns[0]):
raise RuntimeError(f"Unexpected write_file result: {returns[0]!r}")
if len(returns) == 1:
return ModelResponse(
parts=[
ToolCallPart(
tool_name="run_command",
args={
"command": (
f"printf '%s|%s|%s\\n' \"$(cat {STATE_PATH})\" "
'"$AIRFLOW_BOAT_E2E_ENV" "$((6 * 7))"'
)
},
tool_call_id="shell",
)
]
)
shell_out = str(returns[1])
if f"{MARKER}|{ENV_MARKER}|42" not in shell_out:
raise RuntimeError(f"Unexpected run_command result: {shell_out!r}")
if len(returns) == 2:
return ModelResponse(
parts=[
ToolCallPart(tool_name="read_file", args={"path": STATE_PATH}, tool_call_id="read")
]
)
if MARKER not in str(returns[2]):
raise RuntimeError(f"Unexpected read_file result: {returns[2]!r}")
if len(returns) == 3:
return ModelResponse(
parts=[
ToolCallPart(
tool_name="run_command",
args={"command": "echo to-stderr >&2; exit 3"},
tool_call_id="fail",
)
]
)
failed = str(returns[3])
if "[exit code: 3]" not in failed or "to-stderr" not in failed:
raise RuntimeError(f"Unexpected failure result: {failed!r}")
command_cases = [
("echo comment-ok # trailing comment", "comment-ok"),
("cat <<'EOF'\nheredoc-ok\nEOF", "heredoc-ok"),
# Longer than Boat's own 25 s deadline for this call, so a background
# process that held the call open would time it out and fail the step.
("sleep 60 & echo background-ok", "background-ok"),
]
for index, (command, expected) in enumerate(command_cases, start=4):
if len(returns) == index:
return ModelResponse(
parts=[
ToolCallPart(
tool_name="run_command",
args={"command": command, "timeout_seconds": 10},
tool_call_id=f"command-{index}",
)
]
)
if expected not in str(returns[index]) or "timed out" in str(returns[index]):
raise RuntimeError(f"Unexpected command result: {returns[index]!r}")
if len(returns) == 7:
return ModelResponse(
parts=[ToolCallPart(tool_name="list_directory", args={"path": "/tmp"}, tool_call_id="ls")]
)
if "airflow_sandbox_e2e" not in str(returns[7]):
raise RuntimeError(f"Unexpected listing: {returns[7]!r}")
# The in-guest deadline stops a command without Boat having to delete the
# sandbox, so the toolset must not say the sandbox was replaced.
if len(returns) == 8:
return ModelResponse(
parts=[
ToolCallPart(
tool_name="run_command",
args={"command": "sleep 30", "timeout_seconds": 2},
tool_call_id="slow",
)
]
)
timed_out = str(returns[8])
if "[timed out after 2s]" not in timed_out:
raise RuntimeError(f"Expected a timeout, got: {timed_out!r}")
if "sandbox was replaced" in timed_out:
raise RuntimeError("Boat should survive an in-guest timeout, but the sandbox was replaced")
# Which means the file from the first step is still there, and the sandbox
# still runs commands.
if len(returns) == 9:
return ModelResponse(
parts=[
ToolCallPart(
tool_name="read_file", args={"path": STATE_PATH}, tool_call_id="survived"
)
]
)
if MARKER not in str(returns[9]):
raise RuntimeError(f"Files did not survive the timeout: {returns[9]!r}")
if len(returns) == 10:
return ModelResponse(
parts=[
ToolCallPart(
tool_name="run_command",
args={"command": "echo after-timeout-ok"},
tool_call_id="after-timeout",
)
]
)
if "after-timeout-ok" not in str(returns[10]):
raise RuntimeError(f"The sandbox stopped running commands: {returns[10]!r}")
return ModelResponse(parts=[TextPart(content="sandbox boundary e2e passed")])
backend = RecordingBackend(ttl_seconds=900)
agent = Agent(
FunctionModel(model_function),
instructions="Use the sandbox tools as requested.",
toolsets=[
SandboxToolset(
backend,
spec=SandboxSpec(
block_network=False,
env={"AIRFLOW_BOAT_E2E_ENV": ENV_MARKER},
),
)
],
)
result = agent.run_sync("Run the sandbox boundary system test.")
if result.output != "sandbox boundary e2e passed":
raise RuntimeError(f"Unexpected agent output: {result.output!r}")
# One sandbox for the whole run, since the in-guest timeout never replaces it,
# and the toolset deleted it when the run ended.
if len(backend.created) != 1 or backend.created != backend.destroyed:
raise RuntimeError(f"created {backend.created} but destroyed {backend.destroyed}")
return result.output
run_sandbox_agent()
from tests_common.test_utils.system_tests import get_test_run # noqa: E402