chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,231 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from agents.sandbox.manifest import Manifest
|
||||
from agents.sandbox.runtime_session_manager import SandboxRuntimeSessionManager
|
||||
from agents.sandbox.sandboxes.unix_local import (
|
||||
UnixLocalSandboxSession,
|
||||
UnixLocalSandboxSessionState,
|
||||
)
|
||||
from agents.sandbox.session import (
|
||||
CallbackSink,
|
||||
EventPayloadPolicy,
|
||||
Instrumentation,
|
||||
SandboxSessionEvent,
|
||||
SandboxSessionFinishEvent,
|
||||
)
|
||||
from agents.sandbox.session.sinks import ChainedSink, EventSink
|
||||
from agents.sandbox.snapshot import LocalSnapshot, LocalSnapshotSpec, NoopSnapshotSpec
|
||||
|
||||
|
||||
class _EventSink(EventSink):
|
||||
def __init__(self, *, mode: str, on_error: str = "raise") -> None:
|
||||
self.mode = mode # type: ignore[assignment]
|
||||
self.on_error = on_error # type: ignore[assignment]
|
||||
self.payload_policy = None
|
||||
|
||||
async def handle(self, event: SandboxSessionEvent) -> None: # pragma: no cover
|
||||
_ = event
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def _build_session(tmp_path: Path) -> UnixLocalSandboxSession:
|
||||
state = UnixLocalSandboxSessionState(
|
||||
manifest=Manifest(root=str(tmp_path / "workspace")),
|
||||
snapshot=LocalSnapshot(id="x", base_path=tmp_path),
|
||||
)
|
||||
return UnixLocalSandboxSession.from_state(state)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instrumentation_per_op_policy_overrides_default(tmp_path: Path) -> None:
|
||||
events: list[SandboxSessionEvent] = []
|
||||
session = _build_session(tmp_path)
|
||||
sink = CallbackSink(lambda event, _session: events.append(event), mode="sync")
|
||||
sink.bind(session)
|
||||
instrumentation = Instrumentation(
|
||||
sinks=[sink],
|
||||
payload_policy=EventPayloadPolicy(include_exec_output=False),
|
||||
payload_policy_by_op={"exec": EventPayloadPolicy(include_exec_output=True)},
|
||||
)
|
||||
|
||||
event = SandboxSessionFinishEvent(
|
||||
session_id=uuid.uuid4(),
|
||||
seq=1,
|
||||
op="exec",
|
||||
span_id="span_exec",
|
||||
ok=True,
|
||||
duration_ms=0.0,
|
||||
)
|
||||
event.stdout_bytes = b"hello"
|
||||
event.stderr_bytes = b""
|
||||
|
||||
await instrumentation.emit(event)
|
||||
|
||||
assert isinstance(events[0], SandboxSessionFinishEvent)
|
||||
assert events[0].stdout == "hello"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instrumentation_per_sink_policy_overrides_per_op(tmp_path: Path) -> None:
|
||||
first: list[SandboxSessionEvent] = []
|
||||
second: list[SandboxSessionEvent] = []
|
||||
session = _build_session(tmp_path)
|
||||
sink_a = CallbackSink(lambda event, _session: first.append(event), mode="sync")
|
||||
sink_b = CallbackSink(
|
||||
lambda event, _session: second.append(event),
|
||||
mode="sync",
|
||||
payload_policy=EventPayloadPolicy(include_exec_output=True),
|
||||
)
|
||||
sink_a.bind(session)
|
||||
sink_b.bind(session)
|
||||
|
||||
instrumentation = Instrumentation(
|
||||
sinks=[sink_a, sink_b],
|
||||
payload_policy=EventPayloadPolicy(include_exec_output=False),
|
||||
payload_policy_by_op={"exec": EventPayloadPolicy(include_exec_output=False)},
|
||||
)
|
||||
|
||||
event = SandboxSessionFinishEvent(
|
||||
session_id=uuid.uuid4(),
|
||||
seq=1,
|
||||
op="exec",
|
||||
span_id="span_exec",
|
||||
ok=True,
|
||||
duration_ms=0.0,
|
||||
)
|
||||
event.stdout_bytes = b"hello"
|
||||
event.stderr_bytes = b""
|
||||
|
||||
await instrumentation.emit(event)
|
||||
|
||||
assert isinstance(first[0], SandboxSessionFinishEvent)
|
||||
assert isinstance(second[0], SandboxSessionFinishEvent)
|
||||
assert first[0].stdout is None
|
||||
assert second[0].stdout == "hello"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instrumentation_redacts_raw_exec_bytes_when_output_disabled(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
events: list[SandboxSessionEvent] = []
|
||||
session = _build_session(tmp_path)
|
||||
sink = CallbackSink(lambda event, _session: events.append(event), mode="sync")
|
||||
sink.bind(session)
|
||||
instrumentation = Instrumentation(
|
||||
sinks=[sink],
|
||||
payload_policy=EventPayloadPolicy(include_exec_output=False),
|
||||
)
|
||||
|
||||
event = SandboxSessionFinishEvent(
|
||||
session_id=uuid.uuid4(),
|
||||
seq=1,
|
||||
op="exec",
|
||||
span_id="span_exec",
|
||||
ok=True,
|
||||
duration_ms=0.0,
|
||||
)
|
||||
event.stdout_bytes = b"secret"
|
||||
event.stderr_bytes = b"secret2"
|
||||
|
||||
await instrumentation.emit(event)
|
||||
|
||||
assert isinstance(events[0], SandboxSessionFinishEvent)
|
||||
assert events[0].stdout_bytes is None
|
||||
assert events[0].stderr_bytes is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chained_sink_preserves_completion_order_across_modes() -> None:
|
||||
completed = asyncio.Event()
|
||||
|
||||
class SlowBestEffortSink(_EventSink):
|
||||
async def handle(self, event: SandboxSessionEvent) -> None:
|
||||
_ = event
|
||||
await asyncio.sleep(0)
|
||||
completed.set()
|
||||
|
||||
class AssertAfterSink(_EventSink):
|
||||
async def handle(self, event: SandboxSessionEvent) -> None:
|
||||
_ = event
|
||||
assert completed.is_set(), "later sink ran before earlier sink completed"
|
||||
|
||||
sink_a = SlowBestEffortSink(mode="best_effort", on_error="raise")
|
||||
sink_b = AssertAfterSink(mode="sync", on_error="raise")
|
||||
instrumentation = Instrumentation(sinks=[ChainedSink(sink_a, sink_b)])
|
||||
|
||||
event = SandboxSessionFinishEvent(
|
||||
session_id=uuid.uuid4(),
|
||||
seq=1,
|
||||
op="running",
|
||||
span_id="span_running",
|
||||
ok=True,
|
||||
duration_ms=0.0,
|
||||
)
|
||||
await instrumentation.emit(event)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sink_raise_propagates_to_emit() -> None:
|
||||
class _FailingAsyncSink(_EventSink):
|
||||
async def handle(self, event: SandboxSessionEvent) -> None:
|
||||
_ = event
|
||||
await asyncio.sleep(0)
|
||||
raise RuntimeError("boom")
|
||||
|
||||
instrumentation = Instrumentation(sinks=[_FailingAsyncSink(mode="async", on_error="raise")])
|
||||
event = SandboxSessionFinishEvent(
|
||||
session_id=uuid.uuid4(),
|
||||
seq=1,
|
||||
op="running",
|
||||
span_id="span_running",
|
||||
ok=True,
|
||||
duration_ms=0.0,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
await instrumentation.emit(event)
|
||||
|
||||
|
||||
def test_session_manager_uses_custom_snapshot_spec_without_resolving_default(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
called = False
|
||||
|
||||
def _unexpected_default_resolution() -> LocalSnapshotSpec:
|
||||
nonlocal called
|
||||
called = True
|
||||
raise AssertionError("default snapshot resolution should not run")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"agents.sandbox.runtime_session_manager.resolve_default_local_snapshot_spec",
|
||||
_unexpected_default_resolution,
|
||||
)
|
||||
|
||||
custom = LocalSnapshotSpec(base_path=Path("/tmp/custom-sandbox-snapshots"))
|
||||
resolved = SandboxRuntimeSessionManager._resolve_snapshot_spec(custom)
|
||||
|
||||
assert resolved is custom
|
||||
assert called is False
|
||||
|
||||
|
||||
def test_session_manager_falls_back_to_noop_when_default_snapshot_resolution_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
def _raise_os_error() -> LocalSnapshotSpec:
|
||||
raise OSError("read-only home")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"agents.sandbox.runtime_session_manager.resolve_default_local_snapshot_spec",
|
||||
_raise_os_error,
|
||||
)
|
||||
|
||||
resolved = SandboxRuntimeSessionManager._resolve_snapshot_spec(None)
|
||||
|
||||
assert isinstance(resolved, NoopSnapshotSpec)
|
||||
Reference in New Issue
Block a user