Skip to content

Commit f9c7860

Browse files
authored
fix: enforce sandbox cmd timeout on client request
1 parent 9779915 commit f9c7860

6 files changed

Lines changed: 73 additions & 38 deletions

File tree

agentrun/sandbox/api/__aio_data_async_template.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -413,9 +413,12 @@ async def cmd_async(
413413
"command": command,
414414
"cwd": cwd,
415415
}
416-
if timeout is not None:
417-
data["timeout"] = timeout
418-
return await self.post_async("/processes/cmd", data=data)
416+
timeout_config = (
417+
Config(timeout=timeout) if timeout is not None else None
418+
)
419+
return await self.post_async(
420+
"/processes/cmd", data=data, config=timeout_config
421+
)
419422

420423
async def list_processes_async(self):
421424
return await self.get_async("/processes")

agentrun/sandbox/api/__code_interpreter_data_async_template.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -192,9 +192,12 @@ async def cmd_async(
192192
"command": command,
193193
"cwd": cwd,
194194
}
195-
if timeout is not None:
196-
data["timeout"] = timeout
197-
return await self.post_async("/processes/cmd", data=data)
195+
timeout_config = (
196+
Config(timeout=timeout) if timeout is not None else None
197+
)
198+
return await self.post_async(
199+
"/processes/cmd", data=data, config=timeout_config
200+
)
198201

199202
async def list_processes_async(self):
200203
return await self.get_async("/processes")

agentrun/sandbox/api/aio_data.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -606,9 +606,12 @@ async def cmd_async(
606606
"command": command,
607607
"cwd": cwd,
608608
}
609-
if timeout is not None:
610-
data["timeout"] = timeout
611-
return await self.post_async("/processes/cmd", data=data)
609+
timeout_config = (
610+
Config(timeout=timeout) if timeout is not None else None
611+
)
612+
return await self.post_async(
613+
"/processes/cmd", data=data, config=timeout_config
614+
)
612615

613616
def cmd(
614617
self,
@@ -620,9 +623,12 @@ def cmd(
620623
"command": command,
621624
"cwd": cwd,
622625
}
623-
if timeout is not None:
624-
data["timeout"] = timeout
625-
return self.post("/processes/cmd", data=data)
626+
timeout_config = (
627+
Config(timeout=timeout) if timeout is not None else None
628+
)
629+
return self.post(
630+
"/processes/cmd", data=data, config=timeout_config
631+
)
626632

627633
async def list_processes_async(self):
628634
return await self.get_async("/processes")

agentrun/sandbox/api/code_interpreter_data.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -359,9 +359,12 @@ async def cmd_async(
359359
"command": command,
360360
"cwd": cwd,
361361
}
362-
if timeout is not None:
363-
data["timeout"] = timeout
364-
return await self.post_async("/processes/cmd", data=data)
362+
timeout_config = (
363+
Config(timeout=timeout) if timeout is not None else None
364+
)
365+
return await self.post_async(
366+
"/processes/cmd", data=data, config=timeout_config
367+
)
365368

366369
def cmd(
367370
self,
@@ -373,9 +376,12 @@ def cmd(
373376
"command": command,
374377
"cwd": cwd,
375378
}
376-
if timeout is not None:
377-
data["timeout"] = timeout
378-
return self.post("/processes/cmd", data=data)
379+
timeout_config = (
380+
Config(timeout=timeout) if timeout is not None else None
381+
)
382+
return self.post(
383+
"/processes/cmd", data=data, config=timeout_config
384+
)
379385

380386
async def list_processes_async(self):
381387
return await self.get_async("/processes")

tests/unittests/sandbox/api/test_aio_data.py

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -435,29 +435,37 @@ class TestAioProcesses:
435435

436436
def test_cmd(self, api):
437437
api.cmd("ls", "/home")
438-
api.post.assert_called_once_with(
439-
"/processes/cmd",
440-
data={"command": "ls", "cwd": "/home", "timeout": 30},
441-
)
438+
api.post.assert_called_once()
439+
assert api.post.call_args[0] == ("/processes/cmd",)
440+
call_kwargs = api.post.call_args[1]
441+
assert call_kwargs["data"] == {"command": "ls", "cwd": "/home"}
442+
assert isinstance(call_kwargs["config"], Config)
443+
assert call_kwargs["config"].get_timeout() == 30
442444

443445
def test_cmd_no_timeout(self, api):
444446
api.cmd("ls", "/home", timeout=None)
445-
call_data = api.post.call_args[1]["data"]
447+
call_kwargs = api.post.call_args[1]
448+
call_data = call_kwargs["data"]
446449
assert "timeout" not in call_data
450+
assert call_kwargs["config"] is None
447451

448452
@pytest.mark.asyncio
449453
async def test_cmd_async(self, api):
450454
await api.cmd_async("ls", "/home")
451-
api.post_async.assert_called_once_with(
452-
"/processes/cmd",
453-
data={"command": "ls", "cwd": "/home", "timeout": 30},
454-
)
455+
api.post_async.assert_called_once()
456+
assert api.post_async.call_args[0] == ("/processes/cmd",)
457+
call_kwargs = api.post_async.call_args[1]
458+
assert call_kwargs["data"] == {"command": "ls", "cwd": "/home"}
459+
assert isinstance(call_kwargs["config"], Config)
460+
assert call_kwargs["config"].get_timeout() == 30
455461

456462
@pytest.mark.asyncio
457463
async def test_cmd_async_no_timeout(self, api):
458464
await api.cmd_async("ls", "/home", timeout=None)
459-
call_data = api.post_async.call_args[1]["data"]
465+
call_kwargs = api.post_async.call_args[1]
466+
call_data = call_kwargs["data"]
460467
assert "timeout" not in call_data
468+
assert call_kwargs["config"] is None
461469

462470
def test_list_processes(self, api):
463471
api.list_processes()

tests/unittests/sandbox/api/test_code_interpreter_data.py

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from agentrun.sandbox.api.code_interpreter_data import CodeInterpreterDataAPI
88
from agentrun.sandbox.model import CodeLanguage
9+
from agentrun.utils.config import Config
910

1011

1112
@pytest.fixture
@@ -340,29 +341,37 @@ class TestProcesses:
340341

341342
def test_cmd(self, api):
342343
api.cmd("ls", "/home")
343-
api.post.assert_called_once_with(
344-
"/processes/cmd",
345-
data={"command": "ls", "cwd": "/home", "timeout": 30},
346-
)
344+
api.post.assert_called_once()
345+
assert api.post.call_args[0] == ("/processes/cmd",)
346+
call_kwargs = api.post.call_args[1]
347+
assert call_kwargs["data"] == {"command": "ls", "cwd": "/home"}
348+
assert isinstance(call_kwargs["config"], Config)
349+
assert call_kwargs["config"].get_timeout() == 30
347350

348351
def test_cmd_no_timeout(self, api):
349352
api.cmd("ls", "/home", timeout=None)
350-
call_data = api.post.call_args[1]["data"]
353+
call_kwargs = api.post.call_args[1]
354+
call_data = call_kwargs["data"]
351355
assert "timeout" not in call_data
356+
assert call_kwargs["config"] is None
352357

353358
@pytest.mark.asyncio
354359
async def test_cmd_async(self, api):
355360
await api.cmd_async("ls", "/home")
356-
api.post_async.assert_called_once_with(
357-
"/processes/cmd",
358-
data={"command": "ls", "cwd": "/home", "timeout": 30},
359-
)
361+
api.post_async.assert_called_once()
362+
assert api.post_async.call_args[0] == ("/processes/cmd",)
363+
call_kwargs = api.post_async.call_args[1]
364+
assert call_kwargs["data"] == {"command": "ls", "cwd": "/home"}
365+
assert isinstance(call_kwargs["config"], Config)
366+
assert call_kwargs["config"].get_timeout() == 30
360367

361368
@pytest.mark.asyncio
362369
async def test_cmd_async_no_timeout(self, api):
363370
await api.cmd_async("ls", "/home", timeout=None)
364-
call_data = api.post_async.call_args[1]["data"]
371+
call_kwargs = api.post_async.call_args[1]
372+
call_data = call_kwargs["data"]
365373
assert "timeout" not in call_data
374+
assert call_kwargs["config"] is None
366375

367376
def test_list_processes(self, api):
368377
api.list_processes()

0 commit comments

Comments
 (0)