Kazuhiro Sera
576 lines
20 KiB
Python
576 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import inspect
|
|
import io
|
|
from pathlib import Path
|
|
from types import MappingProxyType
|
|
from typing import Any, cast, get_type_hints
|
|
|
|
import pytest
|
|
from typing_extensions import assert_type
|
|
|
|
from agents import RunConfig, Runner
|
|
from agents.sandbox import ExecResult, Manifest, SandboxAgent
|
|
from agents.sandbox.capabilities import Shell
|
|
from agents.sandbox.files import FileEntry
|
|
from agents.sandbox.session.base_sandbox_session import BaseSandboxSession
|
|
from agents.sandbox.session.pty_types import PtyExecUpdate
|
|
from agents.testing import (
|
|
InvalidSandboxStep,
|
|
SandboxCall,
|
|
SandboxCallMatcherError,
|
|
ScriptedModel,
|
|
ScriptedSandboxSession,
|
|
UnconsumedSandboxSteps,
|
|
UnexpectedSandboxCall,
|
|
assistant_message,
|
|
function_call,
|
|
scripted_sandbox_session,
|
|
)
|
|
from agents.testing.sandbox import ScriptedSandboxSession as CanonicalScriptedSandboxSession
|
|
|
|
|
|
class _CallableBytesIO(io.BytesIO):
|
|
def __call__(self) -> None:
|
|
pass
|
|
|
|
|
|
class _CallableBufferedReader(io.BufferedReader):
|
|
def __call__(self) -> None:
|
|
pass
|
|
|
|
|
|
def test_scripted_sandbox_has_public_result_type() -> None:
|
|
session = scripted_sandbox_session()
|
|
assert_type(session, ScriptedSandboxSession)
|
|
base_session: BaseSandboxSession = session
|
|
|
|
assert ScriptedSandboxSession is CanonicalScriptedSandboxSession
|
|
assert get_type_hints(scripted_sandbox_session)["return"] is ScriptedSandboxSession
|
|
assert inspect.isabstract(ScriptedSandboxSession)
|
|
assert isinstance(session, ScriptedSandboxSession)
|
|
assert base_session is session
|
|
|
|
|
|
def test_scripted_sandbox_exposes_only_configured_scriptable_methods() -> None:
|
|
session = scripted_sandbox_session(
|
|
[{"method": "exec", "result": ExecResult(stdout=b"", stderr=b"", exit_code=0)}]
|
|
)
|
|
|
|
assert isinstance(session, BaseSandboxSession)
|
|
assert hasattr(session, "exec")
|
|
assert not hasattr(session, "read")
|
|
assert not hasattr(session, "apply_patch")
|
|
assert not hasattr(session, "pty_exec_start")
|
|
assert "exec" in dir(session)
|
|
assert "read" not in dir(session)
|
|
assert "apply_patch" not in dir(session)
|
|
assert "pty_exec_start" not in dir(session)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_cleanup_does_not_advertise_pty_termination() -> None:
|
|
session = scripted_sandbox_session()
|
|
|
|
assert not hasattr(session, "pty_terminate_all")
|
|
assert "pty_terminate_all" not in dir(session)
|
|
|
|
await session.aclose()
|
|
|
|
async with scripted_sandbox_session() as context_session:
|
|
assert await context_session.running() is True
|
|
|
|
assert await context_session.running() is False
|
|
|
|
|
|
def test_scripted_sandbox_snapshots_manifest_and_derives_pty_support() -> None:
|
|
manifest = Manifest(root="/configured")
|
|
session = scripted_sandbox_session(
|
|
[{"method": "pty_exec_start", "result": None}],
|
|
manifest=manifest,
|
|
)
|
|
manifest.root = "/mutated"
|
|
|
|
assert session.state.manifest.root == "/configured"
|
|
assert session.supports_pty() is True
|
|
|
|
pty_session = scripted_sandbox_session(
|
|
[
|
|
{"method": "pty_exec_start", "result": None},
|
|
{"method": "pty_write_stdin", "result": None},
|
|
]
|
|
)
|
|
assert pty_session.supports_pty() is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("configured_method", "missing_method"),
|
|
[
|
|
("pty_exec_start", "pty_write_stdin"),
|
|
("pty_write_stdin", "pty_exec_start"),
|
|
],
|
|
)
|
|
async def test_scripted_sandbox_exposes_pty_methods_as_one_capability(
|
|
configured_method: str,
|
|
missing_method: str,
|
|
) -> None:
|
|
session = scripted_sandbox_session([{"method": configured_method, "result": None}])
|
|
|
|
assert session.supports_pty() is True
|
|
assert hasattr(session, "pty_exec_start")
|
|
assert hasattr(session, "pty_write_stdin")
|
|
assert "pty_exec_start" in dir(session)
|
|
assert "pty_write_stdin" in dir(session)
|
|
|
|
with pytest.raises(UnexpectedSandboxCall) as exc_info:
|
|
if missing_method == "pty_exec_start":
|
|
await session.pty_exec_start("pwd")
|
|
else:
|
|
await session.pty_write_stdin(session_id=1, chars="")
|
|
|
|
assert exc_info.value.actual_method == missing_method
|
|
assert exc_info.value.expected_method == configured_method
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_supports_capability_method_inventory() -> None:
|
|
read_result = io.BytesIO(b"contents")
|
|
list_result: list[FileEntry] = []
|
|
pty_start_result = PtyExecUpdate(
|
|
process_id=123,
|
|
output=b"started",
|
|
exit_code=None,
|
|
original_token_count=None,
|
|
)
|
|
pty_write_result = PtyExecUpdate(
|
|
process_id=None,
|
|
output=b"done",
|
|
exit_code=0,
|
|
original_token_count=None,
|
|
)
|
|
session = scripted_sandbox_session(
|
|
[
|
|
{"method": "read", "result": read_result},
|
|
{"method": "write", "result": None},
|
|
{"method": "ls", "result": list_result},
|
|
{"method": "mkdir", "result": None},
|
|
{"method": "rm", "result": None},
|
|
{"method": "apply_patch", "result": "Done!"},
|
|
{"method": "pty_exec_start", "result": pty_start_result},
|
|
{"method": "pty_write_stdin", "result": pty_write_result},
|
|
]
|
|
)
|
|
stream = io.BytesIO(b"payload")
|
|
|
|
returned_read_result = cast(io.BytesIO, await session.read(Path("in.txt")))
|
|
assert returned_read_result is not read_result
|
|
assert returned_read_result.getvalue() == b"contents"
|
|
await session.write(Path("out.txt"), stream)
|
|
assert await session.ls(".") == []
|
|
await session.mkdir("new", parents=True)
|
|
await session.rm("old", recursive=True)
|
|
assert await session.apply_patch({"type": "delete_file", "path": "old.txt"}) == "Done!"
|
|
assert (await session.pty_exec_start("sh", tty=True)).process_id == 123
|
|
assert (await session.pty_write_stdin(session_id=123, chars="exit")).exit_code == 0
|
|
|
|
assert [call.method for call in session.calls] == [
|
|
"read",
|
|
"write",
|
|
"ls",
|
|
"mkdir",
|
|
"rm",
|
|
"apply_patch",
|
|
"pty_exec_start",
|
|
"pty_write_stdin",
|
|
]
|
|
recorded_stream = cast(io.BytesIO, session.calls[1].args[1])
|
|
assert recorded_stream is not stream
|
|
assert recorded_stream.getvalue() == b"payload"
|
|
assert recorded_stream.tell() == 0
|
|
session.assert_complete()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("configured", "original_content"),
|
|
[(io.BytesIO(b"before"), b"before"), (io.StringIO("before"), "before")],
|
|
)
|
|
async def test_scripted_sandbox_snapshots_supported_stream_results(
|
|
configured: io.BytesIO | io.StringIO,
|
|
original_content: bytes | str,
|
|
) -> None:
|
|
configured.seek(2)
|
|
session = scripted_sandbox_session([{"method": "read", "result": configured}])
|
|
configured.seek(0)
|
|
configured.write(b"after" if isinstance(configured, io.BytesIO) else "after")
|
|
configured.close()
|
|
|
|
result = cast(io.BytesIO | io.StringIO, await session.read(Path("input.txt")))
|
|
|
|
assert result is not configured
|
|
assert result.tell() == 2
|
|
assert result.getvalue() == original_content
|
|
session.assert_complete()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_snapshots_supported_stream_call_arguments() -> None:
|
|
session = scripted_sandbox_session([{"method": "write", "result": None}])
|
|
source = io.BytesIO(b"before")
|
|
source.seek(3)
|
|
|
|
await session.write(Path("output.bin"), source)
|
|
source.seek(0)
|
|
source.write(b"after")
|
|
source.close()
|
|
|
|
recorded = cast(io.BytesIO, session.calls[0].args[1])
|
|
assert recorded is not source
|
|
assert recorded.tell() == 3
|
|
assert recorded.getvalue() == b"before"
|
|
recorded.write(b"changed")
|
|
retained = cast(io.BytesIO, session.calls[0].args[1])
|
|
assert retained.tell() == 3
|
|
assert retained.getvalue() == b"before"
|
|
session.assert_complete()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_snapshots_callable_supported_streams() -> None:
|
|
result_source = _CallableBytesIO(b"result")
|
|
call_source = _CallableBytesIO(b"call")
|
|
session = scripted_sandbox_session(
|
|
[
|
|
{"method": "read", "result": result_source},
|
|
{"method": "write", "result": None},
|
|
]
|
|
)
|
|
|
|
result = cast(io.BytesIO, await session.read(Path("input.bin")))
|
|
await session.write(Path("output.bin"), call_source)
|
|
|
|
assert result is not result_source
|
|
assert result.getvalue() == b"result"
|
|
recorded = cast(io.BytesIO, session.calls[1].args[1])
|
|
assert recorded is not call_source
|
|
assert recorded.getvalue() == b"call"
|
|
session.assert_complete()
|
|
|
|
|
|
def test_scripted_sandbox_rejects_unsupported_or_closed_stream_results() -> None:
|
|
with io.BufferedReader(io.BytesIO(b"payload")) as unsupported:
|
|
with pytest.raises(InvalidSandboxStep) as unsupported_info:
|
|
scripted_sandbox_session([{"method": "read", "result": unsupported}])
|
|
assert unsupported_info.value.reason == "invalid_outcome"
|
|
|
|
closed = io.BytesIO(b"payload")
|
|
closed.close()
|
|
with pytest.raises(InvalidSandboxStep) as closed_info:
|
|
scripted_sandbox_session([{"method": "read", "result": closed}])
|
|
assert closed_info.value.reason == "invalid_outcome"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_rejects_unsupported_stream_call_before_commit() -> None:
|
|
session = scripted_sandbox_session([{"method": "write", "result": None}])
|
|
|
|
with io.BufferedReader(io.BytesIO(b"payload")) as unsupported:
|
|
with pytest.raises(TypeError, match="support only io.BytesIO and io.StringIO"):
|
|
await session.write(Path("output.bin"), unsupported)
|
|
|
|
assert session.calls == ()
|
|
assert session.remaining_steps == 1
|
|
|
|
|
|
def test_scripted_sandbox_rejects_callable_unsupported_stream_result() -> None:
|
|
with _CallableBufferedReader(io.BytesIO(b"payload")) as unsupported:
|
|
with pytest.raises(InvalidSandboxStep) as exc_info:
|
|
scripted_sandbox_session([{"method": "read", "result": unsupported}])
|
|
|
|
assert exc_info.value.reason == "invalid_outcome"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"stream_factory",
|
|
[
|
|
lambda: _CallableBytesIO(b"payload"),
|
|
lambda: _CallableBufferedReader(io.BytesIO(b"payload")),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
("field", "reason"),
|
|
[("match", "invalid_matcher"), ("responder", "invalid_outcome")],
|
|
)
|
|
def test_scripted_sandbox_rejects_callable_stream_matchers_and_responders(
|
|
stream_factory: Any,
|
|
field: str,
|
|
reason: str,
|
|
) -> None:
|
|
stream = cast(io.IOBase, stream_factory())
|
|
step: dict[str, Any] = {"method": "exec", field: stream}
|
|
if field == "match":
|
|
step["result"] = ExecResult(stdout=b"", stderr=b"", exit_code=0)
|
|
|
|
try:
|
|
with pytest.raises(InvalidSandboxStep) as exc_info:
|
|
scripted_sandbox_session([step])
|
|
finally:
|
|
stream.close()
|
|
|
|
assert exc_info.value.reason == reason
|
|
assert exc_info.value.input_index == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_rejects_callable_unsupported_stream_call_before_commit() -> None:
|
|
session = scripted_sandbox_session([{"method": "write", "result": None}])
|
|
|
|
with _CallableBufferedReader(io.BytesIO(b"payload")) as unsupported:
|
|
with pytest.raises(TypeError, match="support only io.BytesIO and io.StringIO"):
|
|
await session.write(Path("output.bin"), unsupported)
|
|
|
|
assert session.calls == ()
|
|
assert session.remaining_steps == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_snapshots_static_results_when_queued() -> None:
|
|
configured = ExecResult(stdout=b"before", stderr=b"", exit_code=0)
|
|
session = scripted_sandbox_session([{"method": "exec", "result": configured}])
|
|
configured.stdout = b"after"
|
|
|
|
result = await session.exec("pwd")
|
|
|
|
assert result.stdout == b"before"
|
|
session.assert_complete()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_records_detached_fifo_calls() -> None:
|
|
source_operations: list[dict[str, object]] = [{"type": "delete_file", "path": "before.txt"}]
|
|
session = scripted_sandbox_session(
|
|
[
|
|
{"method": "apply_patch", "result": "Done!"},
|
|
{"method": "exec", "result": ExecResult(stdout=b"ok", stderr=b"", exit_code=0)},
|
|
]
|
|
)
|
|
|
|
assert await session.apply_patch(cast(Any, source_operations)) == "Done!"
|
|
source_operations[0]["path"] = "after.txt"
|
|
result = await session.exec("pwd", shell=False)
|
|
|
|
assert result.stdout == b"ok"
|
|
assert session.remaining_steps == 0
|
|
session.assert_complete()
|
|
assert session.calls[0].method == "apply_patch"
|
|
assert session.calls[0].args[0] == [{"type": "delete_file", "path": "before.txt"}]
|
|
assert session.calls[1] == SandboxCall(
|
|
call_index=1,
|
|
method="exec",
|
|
args=("pwd",),
|
|
kwargs=MappingProxyType({"timeout": None, "shell": False, "user": None}),
|
|
)
|
|
|
|
returned_operations = cast(list[dict[str, object]], session.calls[0].args[0])
|
|
returned_operations[0]["path"] = "mutated.txt"
|
|
assert session.calls[0].args[0] == [{"type": "delete_file", "path": "before.txt"}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_snapshot_failure_does_not_commit_call() -> None:
|
|
class Uncopyable:
|
|
def __deepcopy__(self, memo: dict[int, object]) -> object:
|
|
_ = memo
|
|
raise RuntimeError("cannot snapshot")
|
|
|
|
session = scripted_sandbox_session(
|
|
[{"method": "exec", "result": ExecResult(stdout=b"", stderr=b"", exit_code=0)}]
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="cannot snapshot"):
|
|
await session.exec(cast(Any, Uncopyable()))
|
|
|
|
assert session.calls == ()
|
|
assert session.remaining_steps == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_supports_matchers_responders_and_errors() -> None:
|
|
injected_error = RuntimeError("sandbox unavailable")
|
|
session = scripted_sandbox_session(
|
|
[
|
|
{
|
|
"method": "exec",
|
|
"match": lambda call: call.args == ("pwd",),
|
|
"responder": lambda call: ExecResult(
|
|
stdout=f"call {len(call.args)}".encode(), stderr=b"", exit_code=0
|
|
),
|
|
},
|
|
{"method": "exec", "error": injected_error},
|
|
]
|
|
)
|
|
|
|
assert (await session.exec("pwd")).stdout == b"call 1"
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await session.exec("next")
|
|
assert exc_info.value is injected_error
|
|
session.assert_complete()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_reports_structured_payload_free_failures() -> None:
|
|
mismatch = scripted_sandbox_session(
|
|
[
|
|
{"method": "exec", "result": ExecResult(stdout=b"", stderr=b"", exit_code=0)},
|
|
{"method": "read", "result": None},
|
|
]
|
|
)
|
|
|
|
with pytest.raises(UnexpectedSandboxCall) as mismatch_info:
|
|
await mismatch.read(Path("/secret/payload.txt"))
|
|
mismatch_error = mismatch_info.value
|
|
assert mismatch_error.call_index == 0
|
|
assert "call #1" in str(mismatch_error)
|
|
assert mismatch_error.actual_method == "read"
|
|
assert mismatch_error.expected_method == "exec"
|
|
assert mismatch_error.remaining_steps == 2
|
|
assert mismatch.remaining_steps == 2
|
|
assert "/secret/payload.txt" not in str(mismatch_error)
|
|
await mismatch.exec("retry")
|
|
assert mismatch.remaining_steps == 1
|
|
|
|
rejected = scripted_sandbox_session(
|
|
[
|
|
{
|
|
"method": "exec",
|
|
"match": lambda call: call.args == ("expected",),
|
|
"result": ExecResult(stdout=b"", stderr=b"", exit_code=0),
|
|
}
|
|
]
|
|
)
|
|
with pytest.raises(SandboxCallMatcherError) as rejected_info:
|
|
await rejected.exec("secret command")
|
|
assert rejected_info.value.call_index == 0
|
|
assert "call #1" in str(rejected_info.value)
|
|
assert rejected_info.value.method == "exec"
|
|
assert "secret command" not in str(rejected_info.value)
|
|
assert rejected.remaining_steps == 1
|
|
await rejected.exec("expected")
|
|
rejected.assert_complete()
|
|
|
|
extra = scripted_sandbox_session(
|
|
[{"method": "exec", "result": ExecResult(stdout=b"", stderr=b"", exit_code=0)}]
|
|
)
|
|
await extra.exec("first")
|
|
assert hasattr(extra, "exec")
|
|
with pytest.raises(UnexpectedSandboxCall) as extra_info:
|
|
await extra.exec("second secret")
|
|
assert extra_info.value.call_index == 1
|
|
assert "call #2" in str(extra_info.value)
|
|
assert extra_info.value.expected_method is None
|
|
assert extra_info.value.remaining_steps == 0
|
|
assert "second secret" not in str(extra_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_retains_step_after_matcher_exception() -> None:
|
|
matcher_error = RuntimeError("matcher failed")
|
|
matcher_calls = 0
|
|
|
|
def match(_call: SandboxCall) -> bool:
|
|
nonlocal matcher_calls
|
|
matcher_calls += 1
|
|
if matcher_calls == 1:
|
|
raise matcher_error
|
|
return True
|
|
|
|
session = scripted_sandbox_session(
|
|
[
|
|
{
|
|
"method": "exec",
|
|
"match": match,
|
|
"result": ExecResult(stdout=b"ok", stderr=b"", exit_code=0),
|
|
}
|
|
]
|
|
)
|
|
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
await session.exec("first")
|
|
|
|
assert exc_info.value is matcher_error
|
|
assert session.remaining_steps == 1
|
|
assert (await session.exec("retry")).stdout == b"ok"
|
|
session.assert_complete()
|
|
|
|
|
|
def test_scripted_sandbox_reports_unconsumed_steps() -> None:
|
|
session = scripted_sandbox_session(
|
|
[
|
|
{"method": "exec", "result": ExecResult(stdout=b"", stderr=b"", exit_code=0)},
|
|
{"method": "read", "result": None},
|
|
]
|
|
)
|
|
|
|
with pytest.raises(UnconsumedSandboxSteps) as exc_info:
|
|
session.assert_complete()
|
|
assert exc_info.value.remaining_steps == 2
|
|
assert exc_info.value.pending_methods == ("exec", "read")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("step", "reason"),
|
|
[
|
|
("not a mapping", "invalid_input"),
|
|
({"method": "exec", "matc": lambda _call: True, "result": None}, "invalid_input"),
|
|
({"method": "unknown", "result": None}, "unknown_method"),
|
|
({"method": "exec", "match": "no", "result": None}, "invalid_matcher"),
|
|
({"method": "exec"}, "invalid_outcome"),
|
|
({"method": "exec", "result": None, "error": RuntimeError()}, "invalid_outcome"),
|
|
({"method": "exec", "responder": "no"}, "invalid_outcome"),
|
|
({"method": "exec", "error": "no"}, "invalid_outcome"),
|
|
],
|
|
)
|
|
def test_scripted_sandbox_validates_steps_before_use(step: object, reason: str) -> None:
|
|
with pytest.raises(InvalidSandboxStep) as exc_info:
|
|
scripted_sandbox_session(cast(Any, [step]))
|
|
assert exc_info.value.reason == reason
|
|
assert exc_info.value.input_index == 0
|
|
assert "step #1" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scripted_sandbox_drives_black_box_sandbox_agent_workflow() -> None:
|
|
session = scripted_sandbox_session(
|
|
[
|
|
{
|
|
"method": "exec",
|
|
"match": lambda call: call.args == ("pwd",),
|
|
"result": ExecResult(stdout=b"/workspace\n", stderr=b"", exit_code=0),
|
|
}
|
|
]
|
|
)
|
|
model = ScriptedModel(
|
|
[
|
|
[function_call("exec_command", {"cmd": "pwd"}, call_id="call_1")],
|
|
[assistant_message("The workspace is /workspace.")],
|
|
]
|
|
)
|
|
agent = SandboxAgent(
|
|
name="Test agent",
|
|
model=model,
|
|
capabilities=[Shell()],
|
|
)
|
|
|
|
result = await Runner.run(
|
|
agent,
|
|
"Where am I?",
|
|
run_config=RunConfig(sandbox={"session": session}),
|
|
)
|
|
|
|
assert result.final_output == "The workspace is /workspace."
|
|
assert len(session.calls) == 1
|
|
assert len(model.calls) == 2
|
|
session.assert_complete()
|
|
model.assert_complete()
|