diff --git a/rock/actions/sandbox/request.py b/rock/actions/sandbox/request.py index 7a3658df..7c471eb3 100644 --- a/rock/actions/sandbox/request.py +++ b/rock/actions/sandbox/request.py @@ -7,6 +7,15 @@ class Command(BaseModel): session_type: Literal["bash"] = "bash" command: str | list[str] + timeout: float | None = 1200 + """The timeout for the command. None means no timeout.""" + + env: dict[str, str] | None = None + """Environment variables to pass to the command.""" + + cwd: str | None = None + """The current working directory to run the command in.""" + class CreateBashSessionRequest(BaseModel): session_type: Literal["bash"] = "bash" diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index 5c70f92d..b4ad259a 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -171,6 +171,9 @@ async def execute(self, command: Command) -> CommandResponse: data = { "command": command.command, "sandbox_id": self.sandbox_id, + "timeout": command.timeout, + "cwd": command.cwd, + "env": command.env } try: response = await HttpUtils.post(url, headers, data) diff --git a/tests/integration/sdk/sandbox/test_sdk_client.py b/tests/integration/sdk/sandbox/test_sdk_client.py index 318e3adb..5b1d33a1 100644 --- a/tests/integration/sdk/sandbox/test_sdk_client.py +++ b/tests/integration/sdk/sandbox/test_sdk_client.py @@ -52,7 +52,7 @@ async def test_sandbox_get_status(admin_remote_server): await sandbox.start() assert "Failed to start sandbox" in str(exc_info.value) sandbox.stop() - + @pytest.mark.need_admin @SKIP_IF_NO_DOCKER @pytest.mark.asyncio @@ -72,3 +72,18 @@ async def test_update_mount(sandbox_instance: Sandbox): with pytest.raises(Exception) as exc_info: await sandbox_instance.arun(session="default", cmd="touch /tmp/local_files/test.txt") assert "Read-only file system" in str(exc_info.value) + + +@pytest.mark.need_admin +@SKIP_IF_NO_DOCKER +@pytest.mark.asyncio +async def test_execute(sandbox_instance: Sandbox): + from rock.actions.sandbox.request import Command + + curr_status = await sandbox_instance.get_status() + if curr_status.is_alive: + resp1 = await sandbox_instance.execute(Command(command="pwd", cwd="/root")) + assert resp1.stdout.strip() == "/root" + resp2 = await sandbox_instance.execute(Command(command="pwd", cwd="/tmp")) + assert resp2.stdout.strip() == "/tmp" +