from __future__ import annotations import io import uuid from pathlib import Path from typing import Any, cast import pytest from agents.sandbox import Manifest, SandboxPathGrant from agents.sandbox.capabilities import Shell, ShellToolSet from agents.sandbox.capabilities.tools import ( ExecCommandArgs, ExecCommandTool, WriteStdinArgs, WriteStdinTool, ) from agents.sandbox.capabilities.tools.shell_tool import _resolve_shell from agents.sandbox.errors import ExecTimeoutError, ExecTransportError, PtySessionNotFoundError from agents.sandbox.session.base_sandbox_session import BaseSandboxSession from agents.sandbox.session.pty_types import PtyExecUpdate from agents.sandbox.snapshot import NoopSnapshot from agents.sandbox.types import ExecResult, User from agents.tool import FunctionTool from agents.tool_context import ToolContext from tests.utils.factories import TestSessionState class _ShellSession(BaseSandboxSession): def __init__(self, manifest: Manifest) -> None: self.state = TestSessionState( manifest=manifest, snapshot=NoopSnapshot(id=str(uuid.uuid4())), ) self.exec_calls: list[tuple[str, float | None, bool | list[str]]] = [] self.exec_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: object = None) -> io.BytesIO: _ = (path, user) raise AssertionError("read() should not be called") async def write(self, path: Path, data: io.IOBase, *, user: object = None) -> None: _ = (path, data, user) raise AssertionError("write() should not be called") async def _exec_internal( self, *command: str | Path, timeout: float | None = None, ) -> ExecResult: _ = command _ = timeout raise AssertionError("_exec_internal() should not be called directly") async def exec( self, *command: str | Path, timeout: float | None = None, user: str | User | None = None, shell: bool | list[str] = False, ) -> ExecResult: self.exec_users.append(user.name if isinstance(user, User) else user) rendered_command = " ".join(str(part) for part in command) self.exec_calls.append((rendered_command, timeout, shell)) return ExecResult( stdout=f"stdout: {rendered_command}".encode(), stderr=f"stderr: {rendered_command}".encode(), exit_code=7, ) async def persist_workspace(self) -> io.IOBase: return io.BytesIO() async def hydrate_workspace(self, data: io.IOBase) -> None: _ = data class _TimeoutShellSession(_ShellSession): async def exec( self, *command: str | Path, timeout: float | None = None, user: str | User | None = None, shell: bool | list[str] = False, ) -> ExecResult: _ = (command, user, shell) raise ExecTimeoutError(command=("sleep 30",), timeout_s=timeout) class _OutputShellSession(_ShellSession): def __init__( self, manifest: Manifest, *, stdout: bytes, stderr: bytes, exit_code: int = 7, ) -> None: super().__init__(manifest) self.stdout = stdout self.stderr = stderr self.exit_code = exit_code async def exec( self, *command: str | Path, timeout: float | None = None, user: str | User | None = None, shell: bool | list[str] = False, ) -> ExecResult: self.exec_users.append(user.name if isinstance(user, User) else user) rendered_command = " ".join(str(part) for part in command) self.exec_calls.append((rendered_command, timeout, shell)) return ExecResult(stdout=self.stdout, stderr=self.stderr, exit_code=self.exit_code) class _PtyShellSession(_ShellSession): def __init__(self, manifest: Manifest) -> None: super().__init__(manifest) self._next_session_id = 1337 self._live_sessions: set[int] = set() self.last_exec_yield_time_s: float | None = None self.last_exec_user: str | None = None self.last_write_yield_time_s: float | None = None def supports_pty(self) -> bool: return True async def pty_exec_start( self, *command: str | Path, timeout: float | None = None, shell: bool | list[str] = True, user: str | User | None = None, tty: bool = False, yield_time_s: float | None = None, max_output_tokens: int | None = None, ) -> PtyExecUpdate: _ = (command, timeout, shell, tty, max_output_tokens) self.last_exec_user = user.name if isinstance(user, User) else user self.last_exec_yield_time_s = yield_time_s session_id = self._next_session_id self._next_session_id += 1 self._live_sessions.add(session_id) return PtyExecUpdate( process_id=session_id, output=b"", exit_code=None, original_token_count=None, ) async def pty_write_stdin( self, *, session_id: int, chars: str, yield_time_s: float | None = None, max_output_tokens: int | None = None, ) -> PtyExecUpdate: _ = max_output_tokens self.last_write_yield_time_s = yield_time_s if session_id not in self._live_sessions: raise PtySessionNotFoundError(session_id=session_id) self._live_sessions.discard(session_id) return PtyExecUpdate( process_id=None, output=chars.encode("utf-8", errors="replace"), exit_code=0, original_token_count=None, ) class _PtyNoStdinShellSession(_PtyShellSession): async def pty_write_stdin( self, *, session_id: int, chars: str, yield_time_s: float | None = None, max_output_tokens: int | None = None, ) -> PtyExecUpdate: _ = (chars, yield_time_s, max_output_tokens) if session_id not in self._live_sessions: raise PtySessionNotFoundError(session_id=session_id) raise RuntimeError("stdin is not available for this process") class _PtyUnexpectedStdinErrorShellSession(_PtyShellSession): async def pty_write_stdin( self, *, session_id: int, chars: str, yield_time_s: float | None = None, max_output_tokens: int | None = None, ) -> PtyExecUpdate: _ = (session_id, chars, yield_time_s, max_output_tokens) raise RuntimeError("unexpected stdin failure") class _PtyTransportFailingShellSession(_OutputShellSession): def __init__( self, manifest: Manifest, *, stdout: bytes = b"", stderr: bytes = b"", exit_code: int = 0, transport_context: dict[str, object] | None = None, ) -> None: super().__init__(manifest, stdout=stdout, stderr=stderr, exit_code=exit_code) self.transport_context = transport_context or {} self.exec_call_count = 0 def supports_pty(self) -> bool: return True async def exec( self, *command: str | Path, timeout: float | None = None, user: str | User | None = None, shell: bool | list[str] = False, ) -> ExecResult: self.exec_call_count += 1 return await super().exec(*command, timeout=timeout, user=user, shell=shell) async def pty_exec_start( self, *command: str | Path, timeout: float | None = None, shell: bool | list[str] = True, user: str | User | None = None, tty: bool = False, yield_time_s: float | None = None, max_output_tokens: int | None = None, ) -> PtyExecUpdate: _ = (timeout, shell, user, tty, yield_time_s, max_output_tokens) raise ExecTransportError( command=command, context=self.transport_context, cause=RuntimeError("connection closed while reading HTTP status line"), ) def _patch_shell_tool_clock( monkeypatch: pytest.MonkeyPatch, *, chunk_id: str, start: float, end: float, ) -> None: monkeypatch.setattr( "agents.sandbox.capabilities.tools.shell_tool.uuid.uuid4", lambda: uuid.UUID(chunk_id), ) times = iter([start, end]) monkeypatch.setattr( "agents.sandbox.capabilities.tools.shell_tool.time.perf_counter", lambda: next(times), ) class TestShellCapability: def test_resolve_shell_uses_plain_sh_when_login_is_false(self) -> None: assert _resolve_shell(None, login=False) == ["sh", "-c"] def test_tools_requires_bound_session(self) -> None: capability = Shell() with pytest.raises(ValueError, match="Shell capability is not bound to a SandboxSession"): capability.tools() def test_tools_exposes_exec_command_function_tool_after_bind(self) -> None: capability = Shell() capability.bind(_ShellSession(Manifest(root="/workspace"))) tools = capability.tools() assert len(tools) == 1 assert isinstance(tools[0], ExecCommandTool) assert isinstance(tools[0], FunctionTool) assert tools[0].name == "exec_command" def test_tools_exposes_write_stdin_for_pty_sessions(self) -> None: capability = Shell() capability.bind(_PtyShellSession(Manifest(root="/workspace"))) tools = capability.tools() assert len(tools) == 2 assert isinstance(tools[0], ExecCommandTool) assert isinstance(tools[1], WriteStdinTool) assert tools[0].name == "exec_command" assert tools[1].name == "write_stdin" def test_configure_tools_can_customize_shell_approvals_after_clone(self) -> None: async def exec_command_needs_approval( _ctx: Any, params: dict[str, Any], _call_id: str ) -> bool: return str(params["cmd"]).startswith("rm ") async def write_stdin_needs_approval( _ctx: Any, params: dict[str, Any], _call_id: str ) -> bool: return str(params["chars"]) == "\u0003" def configure_tools(toolset: ShellToolSet) -> None: toolset.exec_command.needs_approval = exec_command_needs_approval assert toolset.write_stdin is not None toolset.write_stdin.needs_approval = write_stdin_needs_approval capability = Shell(configure_tools=configure_tools).clone() capability.bind(_PtyShellSession(Manifest(root="/workspace"))) tools = capability.tools() exec_command_tool = cast(ExecCommandTool, tools[0]) write_stdin_tool = cast(WriteStdinTool, tools[1]) assert cast(object, exec_command_tool.needs_approval) is exec_command_needs_approval assert cast(object, write_stdin_tool.needs_approval) is write_stdin_needs_approval def test_configure_tools_can_observe_missing_write_stdin_on_non_pty_session(self) -> None: saw_missing_write_stdin = False def configure_tools(toolset: ShellToolSet) -> None: nonlocal saw_missing_write_stdin saw_missing_write_stdin = toolset.write_stdin is None capability = Shell(configure_tools=configure_tools) capability.bind(_ShellSession(Manifest(root="/workspace"))) tools = capability.tools() assert saw_missing_write_stdin is True assert len(tools) == 1 assert isinstance(tools[0], ExecCommandTool) def test_configure_tools_can_replace_exec_command_tool(self) -> None: replacement_exec_command: ExecCommandTool | None = None def configure_tools(toolset: ShellToolSet) -> None: nonlocal replacement_exec_command replacement_exec_command = ExecCommandTool( session=toolset.exec_command.session, needs_approval=True, ) toolset.exec_command = replacement_exec_command capability = Shell(configure_tools=configure_tools) capability.bind(_ShellSession(Manifest(root="/workspace"))) tools = capability.tools() exec_command_tool = cast(ExecCommandTool, tools[0]) assert replacement_exec_command is not None assert exec_command_tool is replacement_exec_command assert exec_command_tool.needs_approval is True @pytest.mark.asyncio async def test_instructions_match_sandbox_shell_guidance(self) -> None: capability = Shell() instructions = await capability.instructions(Manifest(root="/workspace")) assert ( instructions == "When using the shell:\n" "- Use `exec_command` for shell execution.\n" "- If available, use `write_stdin` to interact with or poll running sessions.\n" "- To interrupt a long-running process via `write_stdin`, start it with " "`tty=true` and send Ctrl-C (`\\u0003`).\n" "- Prefer `rg` and `rg --files` for text/file discovery when available.\n" "- Avoid using Python scripts just to print large file chunks." ) @pytest.mark.asyncio async def test_exec_command_tool_runs_commands_with_source_output_format( self, monkeypatch: pytest.MonkeyPatch, ) -> None: capability = Shell() session = _ShellSession(Manifest(root="/workspace")) capability.bind(session) tool = cast(FunctionTool, capability.tools()[0]) uuids = iter([uuid.UUID("12345678123456781234567812345678")]) times = iter([100.0, 100.25]) monkeypatch.setattr( "agents.sandbox.capabilities.tools.shell_tool.uuid.uuid4", lambda: next(uuids), ) monkeypatch.setattr( "agents.sandbox.capabilities.tools.shell_tool.time.perf_counter", lambda: next(times), ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd", yield_time_ms=1500).model_dump_json(), ) assert session.exec_calls == [("pwd", 1.5, True)] assert ( output == "Chunk ID: 123456\n" "Wall time: 0.2500 seconds\n" "Process exited with code 7\n" "Output:\n" "stdout: pwd\n" "stderr: pwd" ) @pytest.mark.asyncio async def test_exec_command_tool_runs_as_bound_user(self) -> None: capability = Shell() session = _ShellSession(Manifest(root="/workspace")) capability.bind(session) capability.bind_run_as(User(name="sandbox-user")) tool = cast(FunctionTool, capability.tools()[0]) await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd").model_dump_json(), ) assert session.exec_users == ["sandbox-user"] @pytest.mark.asyncio async def test_exec_command_tool_includes_original_token_count_when_truncating( self, monkeypatch: pytest.MonkeyPatch, ) -> None: capability = Shell() session = _ShellSession(Manifest(root="/workspace")) capability.bind(session) tool = cast(FunctionTool, capability.tools()[0]) uuids = iter([uuid.UUID("12345678123456781234567812345678")]) times = iter([200.0, 200.5]) monkeypatch.setattr( "agents.sandbox.capabilities.tools.shell_tool.uuid.uuid4", lambda: next(uuids), ) monkeypatch.setattr( "agents.sandbox.capabilities.tools.shell_tool.time.perf_counter", lambda: next(times), ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd", yield_time_ms=1500, max_output_tokens=2).model_dump_json(), ) assert ( output == "Chunk ID: 123456\n" "Wall time: 0.5000 seconds\n" "Process exited with code 7\n" "Original token count: 6\n" "Output:\n" "Total output lines: 2\n\n" "stdo…4 tokens truncated… pwd" ) @pytest.mark.asyncio async def test_exec_command_tool_wraps_workdir_and_uses_custom_shell( self, monkeypatch: pytest.MonkeyPatch, ) -> None: capability = Shell() session = _ShellSession(Manifest(root="/workspace")) capability.bind(session) tool = cast(FunctionTool, capability.tools()[0]) _patch_shell_tool_clock( monkeypatch, chunk_id="87654321876543218765432187654321", start=300.0, end=300.125, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs( cmd="pwd", workdir="src/project", shell="/bin/bash", login=False, ).model_dump_json(), ) assert session.exec_calls == [ ("cd /workspace/src/project && pwd", 10.0, ["/bin/bash", "-c"]) ] assert ( output == "Chunk ID: 876543\n" "Wall time: 0.1250 seconds\n" "Process exited with code 7\n" "Output:\n" "stdout: cd /workspace/src/project && pwd\n" "stderr: cd /workspace/src/project && pwd" ) @pytest.mark.asyncio async def test_exec_command_tool_allows_extra_path_grant_workdir( self, monkeypatch: pytest.MonkeyPatch, ) -> None: capability = Shell() session = _ShellSession( Manifest( root="/workspace", extra_path_grants=(SandboxPathGrant(path="/tmp", read_only=True),), ) ) capability.bind(session) tool = cast(FunctionTool, capability.tools()[0]) _patch_shell_tool_clock( monkeypatch, chunk_id="11111111111111111111111111111111", start=310.0, end=310.25, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs( cmd="pwd", workdir="/tmp", shell="/bin/bash", login=False, ).model_dump_json(), ) assert session.exec_calls == [("cd /tmp && pwd", 10.0, ["/bin/bash", "-c"])] assert ( output == "Chunk ID: 111111\n" "Wall time: 0.2500 seconds\n" "Process exited with code 7\n" "Output:\n" "stdout: cd /tmp && pwd\n" "stderr: cd /tmp && pwd" ) @pytest.mark.asyncio async def test_exec_command_tool_uses_pty_when_supported( self, monkeypatch: pytest.MonkeyPatch, ) -> None: capability = Shell() session = _PtyShellSession(Manifest(root="/workspace")) capability.bind(session) tool = cast(FunctionTool, capability.tools()[0]) _patch_shell_tool_clock( monkeypatch, chunk_id="abcdef12abcdef12abcdef12abcdef12", start=400.0, end=400.05, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd", yield_time_ms=0, tty=True).model_dump_json(), ) assert session.last_exec_yield_time_s == 0.0 assert ( output == "Chunk ID: abcdef\n" "Wall time: 0.0500 seconds\n" "Process running with session ID 1337\n" "Output:\n" "" ) @pytest.mark.asyncio async def test_exec_command_tool_starts_pty_as_bound_user(self) -> None: capability = Shell() session = _PtyShellSession(Manifest(root="/workspace")) capability.bind(session) capability.bind_run_as(User(name="sandbox-user")) tool = cast(FunctionTool, capability.tools()[0]) await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd", yield_time_ms=0, tty=True).model_dump_json(), ) assert session.last_exec_user == "sandbox-user" @pytest.mark.asyncio async def test_exec_command_tool_formats_timeout_without_exit_code( self, monkeypatch: pytest.MonkeyPatch, ) -> None: capability = Shell() session = _TimeoutShellSession(Manifest(root="/workspace")) capability.bind(session) tool = cast(FunctionTool, capability.tools()[0]) _patch_shell_tool_clock( monkeypatch, chunk_id="fedcba98fedcba98fedcba98fedcba98", start=500.0, end=500.005, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="sleep 30", yield_time_ms=5).model_dump_json(), ) assert ( output == "Chunk ID: fedcba\n" "Wall time: 0.0050 seconds\n" "Output:\n" "Command timed out after 0.005 seconds." ) @pytest.mark.asyncio async def test_exec_command_tool_falls_back_to_one_shot_exec_after_startup_transport_error( self, monkeypatch: pytest.MonkeyPatch, ) -> None: tool = ExecCommandTool( session=_PtyTransportFailingShellSession( Manifest(root="/workspace"), stdout=b"fallback ok", transport_context={"stage": "open_pipe", "retry_safe": True}, ) ) _patch_shell_tool_clock( monkeypatch, chunk_id="44444444444444444444444444444444", start=510.0, end=510.1, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd").model_dump_json(), ) assert "PTY transport failed before the interactive session opened" in output assert "Process exited with code 0" in output assert "Process running with session ID" not in output assert "fallback ok" in output @pytest.mark.asyncio async def test_exec_command_tool_does_not_fall_back_for_tty_sessions(self) -> None: tool = ExecCommandTool( session=_PtyTransportFailingShellSession( Manifest(root="/workspace"), transport_context={"stage": "open_pipe", "retry_safe": True, "tty": True}, ) ) with pytest.raises(ExecTransportError): await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd", tty=True).model_dump_json(), ) @pytest.mark.asyncio async def test_exec_command_tool_does_not_fall_back_for_non_retry_safe_transport_errors( self, ) -> None: tool = ExecCommandTool( session=_PtyTransportFailingShellSession( Manifest(root="/workspace"), transport_context={"stage": "open_pipe"}, ) ) with pytest.raises(ExecTransportError): await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd").model_dump_json(), ) @pytest.mark.asyncio async def test_exec_command_tool_uses_stdout_only_when_stderr_is_empty( self, monkeypatch: pytest.MonkeyPatch, ) -> None: tool = ExecCommandTool( session=_OutputShellSession( Manifest(root="/workspace"), stdout=b"stdout only\n", stderr=b"", ) ) _patch_shell_tool_clock( monkeypatch, chunk_id="11111111111111111111111111111111", start=600.0, end=600.1, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd").model_dump_json(), ) assert ( output == "Chunk ID: 111111\n" "Wall time: 0.1000 seconds\n" "Process exited with code 7\n" "Output:\n" "stdout only\n" ) @pytest.mark.asyncio async def test_exec_command_tool_uses_stderr_only_when_stdout_is_empty( self, monkeypatch: pytest.MonkeyPatch, ) -> None: tool = ExecCommandTool( session=_OutputShellSession( Manifest(root="/workspace"), stdout=b"", stderr=b"stderr only\n", ) ) _patch_shell_tool_clock( monkeypatch, chunk_id="22222222222222222222222222222222", start=700.0, end=700.1, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd").model_dump_json(), ) assert ( output == "Chunk ID: 222222\n" "Wall time: 0.1000 seconds\n" "Process exited with code 7\n" "Output:\n" "stderr only\n" ) @pytest.mark.asyncio async def test_exec_command_tool_does_not_insert_extra_newline_when_stdout_already_has_one( self, monkeypatch: pytest.MonkeyPatch, ) -> None: tool = ExecCommandTool( session=_OutputShellSession( Manifest(root="/workspace"), stdout=b"stdout line\n", stderr=b"stderr line\n", ) ) _patch_shell_tool_clock( monkeypatch, chunk_id="33333333333333333333333333333333", start=800.0, end=800.1, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), ExecCommandArgs(cmd="pwd").model_dump_json(), ) assert ( output == "Chunk ID: 333333\n" "Wall time: 0.1000 seconds\n" "Process exited with code 7\n" "Output:\n" "stdout line\n" "stderr line\n" ) @pytest.mark.asyncio async def test_write_stdin_tool_writes_and_finishes_session( self, monkeypatch: pytest.MonkeyPatch, ) -> None: session = _PtyShellSession(Manifest(root="/workspace")) session._live_sessions.add(1337) tool = WriteStdinTool(session=session) _patch_shell_tool_clock( monkeypatch, chunk_id="55555555555555555555555555555555", start=900.0, end=900.2, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), WriteStdinArgs(session_id=1337, chars="hello").model_dump_json(), ) assert ( output == "Chunk ID: 555555\n" "Wall time: 0.2000 seconds\n" "Process exited with code 0\n" "Output:\n" "hello" ) @pytest.mark.asyncio async def test_write_stdin_tool_rejects_non_pty_sessions(self) -> None: tool = WriteStdinTool(session=_ShellSession(Manifest(root="/workspace"))) with pytest.raises( RuntimeError, match="write_stdin is not available for non-PTY sandboxes" ): await tool.on_invoke_tool( cast(ToolContext[object], None), WriteStdinArgs(session_id=1337).model_dump_json(), ) @pytest.mark.asyncio async def test_write_stdin_tool_formats_unknown_session_error( self, monkeypatch: pytest.MonkeyPatch, ) -> None: tool = WriteStdinTool(session=_PtyShellSession(Manifest(root="/workspace"))) _patch_shell_tool_clock( monkeypatch, chunk_id="66666666666666666666666666666666", start=910.0, end=910.1, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), WriteStdinArgs(session_id=9999).model_dump_json(), ) assert ( output == "Chunk ID: 666666\n" "Wall time: 0.1000 seconds\n" "Process exited with code 1\n" "Output:\n" "write_stdin failed: PTY session not found: 9999" ) @pytest.mark.asyncio async def test_write_stdin_tool_formats_missing_stdin_error( self, monkeypatch: pytest.MonkeyPatch, ) -> None: session = _PtyNoStdinShellSession(Manifest(root="/workspace")) session._live_sessions.add(1337) tool = WriteStdinTool(session=session) _patch_shell_tool_clock( monkeypatch, chunk_id="77777777777777777777777777777777", start=920.0, end=920.05, ) output = await tool.on_invoke_tool( cast(ToolContext[object], None), WriteStdinArgs(session_id=1337).model_dump_json(), ) assert ( output == "Chunk ID: 777777\n" "Wall time: 0.0500 seconds\n" "Process exited with code 1\n" "Output:\n" "stdin is not available for this process. Start the command with `tty=true` in " "`exec_command` before using `write_stdin`." ) @pytest.mark.asyncio async def test_write_stdin_tool_reraises_unexpected_runtime_error(self) -> None: tool = WriteStdinTool( session=_PtyUnexpectedStdinErrorShellSession(Manifest(root="/workspace")) ) with pytest.raises(RuntimeError, match="unexpected stdin failure"): await tool.on_invoke_tool( cast(ToolContext[object], None), WriteStdinArgs(session_id=1337).model_dump_json(), )