Skip to content

Commit 4a8a0ab

Browse files
committed
feat(cli): Add exec command
Running a command in a sandbox requires a Python script that creates a Process, drains its output, and handles the exit code. Add `aviato exec` to run a command and stream its stdout/stderr to the terminal. Supports --cwd and --timeout. Exits with the command's return code.
1 parent 17a3397 commit 4a8a0ab

File tree

8 files changed

+286
-3
lines changed

8 files changed

+286
-3
lines changed

AGENTS.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,9 +296,12 @@ Commands:
296296

297297
| Command | File | Description |
298298
|---------|------|-------------|
299+
| `aviato exec` | `cli/exec.py` | Execute a command in a sandbox (`--cwd`, `--timeout`) |
299300
| `aviato list` | `cli/list.py` | List sandboxes with optional filters (`--status`, `--tag`, `--runway-id`, `--tower-id`) |
300301

301302
```bash
303+
aviato exec <sandbox-id> echo hello # Run a command
304+
aviato exec <sandbox-id> --cwd /app ls -la # Run with working directory
302305
aviato list # List all sandboxes
303306
aviato list --status running --tag my-project # Filter by status and tag
304307
```

DEVELOPMENT.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ aviato --help # Show available commands
150150
aviato --version # Show SDK version
151151
aviato list # List sandboxes
152152
aviato list --status running --tag dev # Filter sandboxes
153+
aviato exec <sandbox-id> echo hello # Run a command in a sandbox
153154
```
154155

155156
You can also invoke it as a Python module:

docs/guides/execution.md

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,18 @@ for p in pending:
301301
print(p.result().stdout)
302302
```
303303

304+
## CLI
305+
306+
The `aviato exec` command runs a one-off command in a sandbox from the terminal:
307+
308+
```bash
309+
aviato exec <sandbox-id> echo hello
310+
aviato exec <sandbox-id> --cwd /app ls -la
311+
aviato exec <sandbox-id> --timeout 30 python script.py
312+
```
313+
314+
Exits with the command's return code.
315+
304316
## Process Control
305317

306318
The `Process` object provides methods for monitoring and control:

src/aviato/cli/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
) from e
2222
raise
2323

24+
from aviato.cli.exec import exec_command
2425
from aviato.cli.list import list_sandboxes
2526
from aviato.exceptions import AviatoError
2627

@@ -45,4 +46,5 @@ def cli() -> None:
4546
"""Aviato sandbox CLI."""
4647

4748

49+
cli.add_command(exec_command, "exec")
4850
cli.add_command(list_sandboxes)

src/aviato/cli/exec.py

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
# SPDX-FileCopyrightText: 2025 CoreWeave, Inc.
2+
# SPDX-License-Identifier: Apache-2.0
3+
# SPDX-PackageName: aviato-client
4+
5+
"""aviato exec — execute a command in a sandbox."""
6+
7+
from __future__ import annotations
8+
9+
import sys
10+
import threading
11+
12+
import click
13+
14+
from aviato import Sandbox
15+
16+
17+
@click.command(context_settings={"ignore_unknown_options": True})
18+
@click.argument("sandbox_id")
19+
@click.argument("command", nargs=-1, required=True, type=click.UNPROCESSED)
20+
@click.option(
21+
"--cwd",
22+
default=None,
23+
help="Working directory for the command.",
24+
)
25+
@click.option(
26+
"--timeout",
27+
"timeout_seconds",
28+
type=click.FloatRange(min=0),
29+
default=None,
30+
help="Timeout in seconds.",
31+
)
32+
def exec_command(
33+
sandbox_id: str,
34+
command: tuple[str, ...],
35+
cwd: str | None,
36+
timeout_seconds: float | None,
37+
) -> None:
38+
"""Execute a command in a sandbox.
39+
40+
SANDBOX_ID is the ID of the sandbox to run the command in.
41+
42+
Examples:
43+
44+
aviato exec <sandbox-id> echo hello
45+
46+
aviato exec <sandbox-id> sh -c "echo hello > /tmp/file"
47+
48+
aviato exec <sandbox-id> --cwd /app python script.py
49+
"""
50+
sandbox = Sandbox.from_id(sandbox_id).result()
51+
52+
try:
53+
process = sandbox.exec(
54+
list(command),
55+
cwd=cwd,
56+
timeout_seconds=timeout_seconds,
57+
)
58+
59+
def _drain_stderr() -> None:
60+
try:
61+
for chunk in process.stderr:
62+
click.echo(chunk, nl=False, err=True)
63+
except BrokenPipeError:
64+
pass
65+
66+
stderr_thread = threading.Thread(target=_drain_stderr)
67+
stderr_thread.start()
68+
69+
for line in process.stdout:
70+
click.echo(line, nl=False)
71+
72+
result = process.result()
73+
stderr_thread.join()
74+
except KeyboardInterrupt:
75+
sys.exit(130)
76+
except BrokenPipeError:
77+
pass # Piped to head/etc — exit cleanly
78+
else:
79+
sys.exit(result.returncode)

tests/AGENTS.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ Tests skip gracefully with clear messages when no auth is configured. The `requi
111111
| `test_utilities.py` | aviato.result(), aviato.wait() utilities |
112112
| `test_cli_main.py` | `__main__` entry point, --help, --version, ImportError fallback |
113113
| `test_cli_list.py` | CLI list command, filters, empty state, API errors |
114+
| `test_cli_exec.py` | CLI exec command, stdout, stderr, returncode, cwd, timeout, not-found error |
114115

115116
### Integration Test Files (`tests/integration/aviato/`)
116117

tests/unit/aviato/conftest.py

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@
1818

1919
T = TypeVar("T")
2020

21+
# Event loops created by _make_reader, closed after each test.
22+
_test_loops: list[asyncio.AbstractEventLoop] = []
23+
2124
# Environment variables that affect authentication behavior.
2225
# These are cleared before each test to ensure isolation.
2326
AUTH_ENV_VARS = (
@@ -42,6 +45,14 @@ def clean_auth_env(monkeypatch: pytest.MonkeyPatch) -> None:
4245
monkeypatch.delenv(var, raising=False)
4346

4447

48+
@pytest.fixture(autouse=True)
49+
def _cleanup_test_loops() -> None:
50+
"""Close event loops created by _make_reader after each test."""
51+
yield # type: ignore[misc]
52+
while _test_loops:
53+
_test_loops.pop().close()
54+
55+
4556
@pytest.fixture
4657
def mock_aviato_api_key(monkeypatch: pytest.MonkeyPatch) -> str:
4758
"""Set a mock AVIATO_API_KEY for the test."""
@@ -119,8 +130,18 @@ def make_process(
119130
if stderr:
120131
stderr_queue.put_nowait(None) # Sentinel
121132

122-
mock_lm = MagicMock()
123-
stdout_reader = StreamReader(stdout_queue, mock_lm)
124-
stderr_reader = StreamReader(stderr_queue, mock_lm)
133+
# Each StreamReader gets its own mock _LoopManager with a dedicated event
134+
# loop so that run_sync actually executes the asyncio.Queue.get() coroutine.
135+
# Separate loops are needed because the CLI drains stdout and stderr on
136+
# different threads; a shared loop would hit "already running" errors.
137+
def _make_reader(queue: asyncio.Queue[str | None]) -> StreamReader[str]:
138+
loop = asyncio.new_event_loop()
139+
_test_loops.append(loop)
140+
lm = MagicMock()
141+
lm.run_sync.side_effect = loop.run_until_complete
142+
return StreamReader(queue, lm)
143+
144+
stdout_reader = _make_reader(stdout_queue)
145+
stderr_reader = _make_reader(stderr_queue)
125146

126147
return Process(future, command or [], stdout_reader, stderr_reader)

tests/unit/aviato/test_cli_exec.py

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
# SPDX-FileCopyrightText: 2025 CoreWeave, Inc.
2+
# SPDX-License-Identifier: Apache-2.0
3+
# SPDX-PackageName: aviato-client
4+
5+
"""Tests for aviato exec CLI command."""
6+
7+
from __future__ import annotations
8+
9+
from unittest.mock import MagicMock, patch
10+
11+
from click.testing import CliRunner
12+
13+
from aviato.cli import cli
14+
from aviato.exceptions import SandboxNotFoundError
15+
from tests.unit.aviato.conftest import make_operation_ref, make_process
16+
17+
18+
def _patch_sandbox(process):
19+
"""Patch aviato.cli.exec.Sandbox.from_id to return a sandbox whose exec() returns *process*.
20+
21+
Returns a context manager that patches the Sandbox class.
22+
"""
23+
mock_sandbox = MagicMock()
24+
mock_sandbox.exec.return_value = process
25+
return patch(
26+
"aviato.cli.exec.Sandbox",
27+
**{"from_id.return_value": make_operation_ref(mock_sandbox)},
28+
)
29+
30+
31+
class TestExecCommand:
32+
"""Tests for the aviato exec CLI command."""
33+
34+
def test_exec_prints_stdout(self) -> None:
35+
"""aviato exec prints command stdout and exits with returncode."""
36+
process = make_process(stdout="hello world\n", command=["echo", "hello"])
37+
38+
with _patch_sandbox(process) as mock_cls:
39+
mock_sandbox = mock_cls.from_id.return_value.result()
40+
runner = CliRunner()
41+
result = runner.invoke(cli, ["exec", "test-sandbox-id", "echo", "hello"])
42+
43+
assert result.exit_code == 0
44+
assert "hello world\n" in result.output
45+
mock_sandbox.exec.assert_called_once_with(
46+
["echo", "hello"],
47+
cwd=None,
48+
timeout_seconds=None,
49+
)
50+
51+
def test_exec_nonzero_returncode(self) -> None:
52+
"""aviato exec exits with the process returncode."""
53+
process = make_process(returncode=1, command=["false"])
54+
55+
with _patch_sandbox(process):
56+
runner = CliRunner()
57+
result = runner.invoke(cli, ["exec", "test-sandbox-id", "false"])
58+
59+
assert result.exit_code == 1
60+
61+
def test_exec_with_cwd(self) -> None:
62+
"""aviato exec --cwd passes cwd to sandbox.exec."""
63+
process = make_process(command=["ls"])
64+
65+
with _patch_sandbox(process) as mock_cls:
66+
mock_sandbox = mock_cls.from_id.return_value.result()
67+
runner = CliRunner()
68+
result = runner.invoke(cli, ["exec", "test-sandbox-id", "--cwd", "/app", "ls"])
69+
70+
assert result.exit_code == 0
71+
mock_sandbox.exec.assert_called_once_with(
72+
["ls"],
73+
cwd="/app",
74+
timeout_seconds=None,
75+
)
76+
77+
def test_exec_with_timeout(self) -> None:
78+
"""aviato exec --timeout passes timeout_seconds to sandbox.exec."""
79+
process = make_process(command=["sleep", "10"])
80+
81+
with _patch_sandbox(process) as mock_cls:
82+
mock_sandbox = mock_cls.from_id.return_value.result()
83+
runner = CliRunner()
84+
result = runner.invoke(
85+
cli, ["exec", "test-sandbox-id", "--timeout", "30", "sleep", "10"]
86+
)
87+
88+
assert result.exit_code == 0
89+
mock_sandbox.exec.assert_called_once_with(
90+
["sleep", "10"],
91+
cwd=None,
92+
timeout_seconds=30.0,
93+
)
94+
95+
def test_exec_concurrent_stdout_stderr(self) -> None:
96+
"""aviato exec drains stdout and stderr concurrently on separate threads."""
97+
process = make_process(
98+
stdout="out line 1\nout line 2\n",
99+
stderr="err line 1\nerr line 2\n",
100+
returncode=0,
101+
command=["some-cmd"],
102+
)
103+
104+
with _patch_sandbox(process):
105+
runner = CliRunner(mix_stderr=False)
106+
result = runner.invoke(cli, ["exec", "test-sandbox-id", "some-cmd"])
107+
108+
assert result.exit_code == 0
109+
assert "out line 1" in result.output
110+
assert "out line 2" in result.output
111+
assert "err line 1" in result.stderr
112+
assert "err line 2" in result.stderr
113+
114+
def test_exec_prints_stderr(self) -> None:
115+
"""aviato exec prints command stderr to stderr."""
116+
process = make_process(
117+
stderr="error: something went wrong\n",
118+
returncode=1,
119+
command=["bad-cmd"],
120+
)
121+
122+
with _patch_sandbox(process):
123+
runner = CliRunner(mix_stderr=False)
124+
result = runner.invoke(cli, ["exec", "test-sandbox-id", "bad-cmd"])
125+
126+
assert result.exit_code == 1
127+
assert "error: something went wrong" in result.stderr
128+
129+
def test_exec_sandbox_not_found(self) -> None:
130+
"""aviato exec shows clean error for SandboxNotFoundError."""
131+
mock_op_ref = MagicMock()
132+
mock_op_ref.result.side_effect = SandboxNotFoundError("not found", sandbox_id="bad-id")
133+
134+
with patch("aviato.cli.exec.Sandbox") as mock_sandbox_cls:
135+
mock_sandbox_cls.from_id.return_value = mock_op_ref
136+
137+
runner = CliRunner()
138+
result = runner.invoke(cli, ["exec", "bad-id", "echo", "hello"])
139+
140+
assert result.exit_code == 1
141+
assert "not found" in result.output
142+
143+
def test_exec_keyboard_interrupt(self) -> None:
144+
"""aviato exec exits 130 on KeyboardInterrupt."""
145+
mock_process = MagicMock()
146+
mock_process.stdout = iter([])
147+
mock_process.result.side_effect = KeyboardInterrupt
148+
149+
with _patch_sandbox(mock_process):
150+
runner = CliRunner()
151+
result = runner.invoke(cli, ["exec", "test-sandbox-id", "sleep", "100"])
152+
153+
assert result.exit_code == 130
154+
155+
def test_exec_broken_pipe(self) -> None:
156+
"""aviato exec exits 0 on BrokenPipeError (piped to head/etc)."""
157+
mock_process = MagicMock()
158+
mock_process.stdout.__iter__ = MagicMock(side_effect=BrokenPipeError)
159+
160+
with _patch_sandbox(mock_process):
161+
runner = CliRunner()
162+
result = runner.invoke(cli, ["exec", "test-sandbox-id", "echo", "hello"])
163+
164+
assert result.exit_code == 0

0 commit comments

Comments
 (0)