wehub-resource-sync
201 lines
7.0 KiB
Python
201 lines
7.0 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import cast
|
|
|
|
import pytest
|
|
|
|
from agents.sandbox import Manifest
|
|
from agents.sandbox.capabilities.tools import ViewImageTool
|
|
from agents.sandbox.errors import WorkspaceReadNotFoundError
|
|
from agents.sandbox.session.base_sandbox_session import BaseSandboxSession
|
|
from agents.sandbox.snapshot import NoopSnapshot
|
|
from agents.sandbox.types import ExecResult, User
|
|
from agents.tool import ToolOutputImage
|
|
from agents.tool_context import ToolContext
|
|
from tests.utils.factories import TestSessionState
|
|
|
|
_MAX_IMAGE_BYTES = 10 * 1024 * 1024
|
|
_PNG_BASE64 = (
|
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+a84QAAAAASUVORK5CYII="
|
|
)
|
|
_PNG_BYTES = base64.b64decode(_PNG_BASE64)
|
|
|
|
|
|
class _ImageSession(BaseSandboxSession):
|
|
def __init__(self, manifest: Manifest) -> None:
|
|
self.state = TestSessionState(
|
|
manifest=manifest,
|
|
snapshot=NoopSnapshot(id=str(uuid.uuid4())),
|
|
)
|
|
self.files: dict[Path, bytes] = {}
|
|
self.read_users: list[str | None] = []
|
|
|
|
async def start(self) -> None:
|
|
return None
|
|
|
|
async def stop(self) -> None:
|
|
return None
|
|
|
|
async def shutdown(self) -> None:
|
|
return None
|
|
|
|
async def running(self) -> bool:
|
|
return True
|
|
|
|
async def read(self, path: Path, *, user: str | User | None = None) -> io.BytesIO:
|
|
self.read_users.append(user.name if isinstance(user, User) else user)
|
|
normalized = self.normalize_path(path)
|
|
if normalized not in self.files:
|
|
raise FileNotFoundError(normalized)
|
|
return io.BytesIO(self.files[normalized])
|
|
|
|
async def write(
|
|
self,
|
|
path: Path,
|
|
data: io.IOBase,
|
|
*,
|
|
user: str | User | None = None,
|
|
) -> None:
|
|
_ = user
|
|
normalized = self.normalize_path(path)
|
|
payload = data.read()
|
|
if isinstance(payload, str):
|
|
self.files[normalized] = payload.encode("utf-8")
|
|
else:
|
|
self.files[normalized] = bytes(payload)
|
|
|
|
async def _exec_internal(
|
|
self,
|
|
*command: str | Path,
|
|
timeout: float | None = None,
|
|
) -> ExecResult:
|
|
_ = (command, timeout)
|
|
raise AssertionError("_exec_internal() should not be called")
|
|
|
|
async def persist_workspace(self) -> io.IOBase:
|
|
return io.BytesIO()
|
|
|
|
async def hydrate_workspace(self, data: io.IOBase) -> None:
|
|
_ = data
|
|
|
|
|
|
class _ProviderNotFoundImageSession(_ImageSession):
|
|
async def read(self, path: Path, *, user: str | User | None = None) -> io.BytesIO:
|
|
self.read_users.append(user.name if isinstance(user, User) else user)
|
|
normalized = self.normalize_path(path)
|
|
if normalized in self.files:
|
|
return io.BytesIO(self.files[normalized])
|
|
raise WorkspaceReadNotFoundError(path=normalized)
|
|
|
|
|
|
class TestViewImageTool:
|
|
def test_view_image_accepts_needs_approval_setting(self) -> None:
|
|
session = _ImageSession(Manifest(root="/workspace"))
|
|
|
|
async def needs_approval(_ctx: object, params: dict[str, object], _call_id: str) -> bool:
|
|
return str(params["path"]).startswith("sensitive/")
|
|
|
|
tool = ViewImageTool(session=session, needs_approval=needs_approval)
|
|
|
|
assert cast(object, tool.needs_approval) is needs_approval
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_view_image_returns_tool_output_image_for_png(self) -> None:
|
|
session = _ImageSession(Manifest(root="/workspace"))
|
|
session.files[Path("/workspace/images/dot.png")] = _PNG_BYTES
|
|
tool = ViewImageTool(session=session)
|
|
|
|
output = await tool.on_invoke_tool(
|
|
cast(ToolContext[object], None),
|
|
'{"path":"images/dot.png"}',
|
|
)
|
|
|
|
assert isinstance(output, ToolOutputImage)
|
|
assert output.image_url == f"data:image/png;base64,{_PNG_BASE64}"
|
|
assert output.detail is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_view_image_reads_as_bound_user(self) -> None:
|
|
session = _ImageSession(Manifest(root="/workspace"))
|
|
session.files[Path("/workspace/images/dot.png")] = _PNG_BYTES
|
|
tool = ViewImageTool(session=session, user=User(name="sandbox-user"))
|
|
|
|
output = await tool.on_invoke_tool(
|
|
cast(ToolContext[object], None),
|
|
'{"path":"images/dot.png"}',
|
|
)
|
|
|
|
assert isinstance(output, ToolOutputImage)
|
|
assert session.read_users == ["sandbox-user"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_view_image_rejects_non_image_files(self) -> None:
|
|
session = _ImageSession(Manifest(root="/workspace"))
|
|
session.files[Path("/workspace/notes.txt")] = b"hello\n"
|
|
tool = ViewImageTool(session=session)
|
|
|
|
output = await tool.on_invoke_tool(
|
|
cast(ToolContext[object], None),
|
|
'{"path":"notes.txt"}',
|
|
)
|
|
|
|
assert output == "image path `notes.txt` is not a supported image file"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_view_image_rejects_images_larger_than_10mb(self) -> None:
|
|
session = _ImageSession(Manifest(root="/workspace"))
|
|
session.files[Path("/workspace/images/huge.png")] = b"\x89PNG\r\n\x1a\n" + (
|
|
b"0" * (_MAX_IMAGE_BYTES + 1)
|
|
)
|
|
tool = ViewImageTool(session=session)
|
|
|
|
output = await tool.on_invoke_tool(
|
|
cast(ToolContext[object], None),
|
|
'{"path":"images/huge.png"}',
|
|
)
|
|
|
|
assert output == (
|
|
"image path `images/huge.png` exceeded the allowed size of 10MB; "
|
|
"resize or compress the image and try again"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_view_image_rejection_text_does_not_expose_provider_path(self) -> None:
|
|
provider_root = Path("/provider/private/root")
|
|
session = _ProviderNotFoundImageSession(Manifest(root=str(provider_root)))
|
|
session.files[provider_root / "notes.txt"] = b"hello\n"
|
|
session.files[provider_root / "images/huge.png"] = b"\x89PNG\r\n\x1a\n" + (
|
|
b"0" * (_MAX_IMAGE_BYTES + 1)
|
|
)
|
|
tool = ViewImageTool(session=session)
|
|
|
|
missing_output = await tool.on_invoke_tool(
|
|
cast(ToolContext[object], None),
|
|
'{"path":"images/missing.png"}',
|
|
)
|
|
non_image_output = await tool.on_invoke_tool(
|
|
cast(ToolContext[object], None),
|
|
'{"path":"notes.txt"}',
|
|
)
|
|
huge_output = await tool.on_invoke_tool(
|
|
cast(ToolContext[object], None),
|
|
'{"path":"images/huge.png"}',
|
|
)
|
|
|
|
outputs = [missing_output, non_image_output, huge_output]
|
|
assert outputs == [
|
|
"image path `images/missing.png` was not found",
|
|
"image path `notes.txt` is not a supported image file",
|
|
(
|
|
"image path `images/huge.png` exceeded the allowed size of 10MB; "
|
|
"resize or compress the image and try again"
|
|
),
|
|
]
|
|
for output in outputs:
|
|
assert isinstance(output, str)
|
|
assert str(provider_root) not in output
|