diff --git a/.changeset/exec-session.md b/.changeset/exec-session.md new file mode 100644 index 0000000..dbcec66 --- /dev/null +++ b/.changeset/exec-session.md @@ -0,0 +1,23 @@ +--- +"@upstash/box": patch +--- + +Add `box.exec.session()` — a live, interactive command session over a WebSocket. + +Unlike `exec.command` / `exec.stream`, which run a command and hand back its +result, a session returns a handle to a *running* process: + +- `argv` runs a program directly (no shell); `cmd` runs one via `bash -lc`. +- `write()` sends stdin, `endStdin()` closes it so a command reading to EOF can + finish, and `onStdout` / `onStderr` receive output as it arrives (separate + streams unless `tty` is set). +- `tty` allocates a real PTY sized by `rows`/`cols`, with `resize()` for later + changes — enough for interactive programs and terminal UIs. +- `kill(signal)` sends an allowlisted signal; `terminate(graceMs)` asks the + server for SIGTERM then SIGKILL after the grace. +- `wait()` resolves with the exit code; `close()` hangs up, which also stops the + process. + +Node-only: authentication uses a request header, which browsers cannot set on a +WebSocket handshake. `ws` moves from a dev dependency to a runtime dependency; +the public types stay free of `@types/ws`. diff --git a/packages/python-sdk/CHANGELOG.md b/packages/python-sdk/CHANGELOG.md index 61a178c..df42296 100644 --- a/packages/python-sdk/CHANGELOG.md +++ b/packages/python-sdk/CHANGELOG.md @@ -4,6 +4,16 @@ All notable changes to `upstash-box` (Python) are documented here. ## 0.3.0 +- `exec.session(...)` — live command sessions over a WebSocket, matching + `@upstash/box`. Returns a handle once the process is running, with `pid`, + `exec_id`, `write`, `end_stdin`, `resize`, `kill`, `terminate`, `wait`, and + `close`. Pass `argv` to run a program without a shell or `cmd` to go through + `bash -lc`, `tty=True` for a PTY (with `rows`/`cols`), plus `cwd` and `env` + overlays. `on_stdout`/`on_stderr` receive `bytes` as they arrive. The handle + owns the process: closing it, or losing the connection, kills the command. + Available on both clients and usable as a context manager. +- Adds a `websockets>=13` dependency, imported lazily so it only loads when a + session is opened. - `files.stat(path, follow=...)`, `files.mkdir(path, parents=...)`, `files.rename(from_path, to_path)`, and `files.remove(path, recursive=...)` — filesystem metadata and mutation operations. `stat` returns the entry type diff --git a/packages/python-sdk/PARITY.md b/packages/python-sdk/PARITY.md index b19dfe8..614eaee 100644 --- a/packages/python-sdk/PARITY.md +++ b/packages/python-sdk/PARITY.md @@ -27,6 +27,7 @@ JS `Run`/`StreamRun` → Python `Run`/`StreamRun` (+ `AsyncRun`/`AsyncStreamRun` | `exec.command` / `code` / `stream` / `streamCode` | `exec.command` / `code` / `stream` / `stream_code` | | `files.read/write/list/upload/download` | `files.read/write/list/upload/download` | | `files.stat/mkdir/rename/remove` | `files.stat/mkdir/rename/remove` | +| `exec.session` (live WebSocket session) | `exec.session` | | `git.clone/diff/status/commit/updateConfig/push/createPR/exec/checkout` | `git.clone/diff/status/commit/update_config/push/create_pr/exec/checkout` | | `schedule.exec/agent/list/get/update/pause/resume/delete` | same (snake) | | `skills.add/remove/list` | `skills.add/remove/list` | @@ -122,3 +123,29 @@ statics `create`, `from_snapshot`, `get_by_name`, `delete_boxes`, `helpers` → covered by `tests/helpers.py`; `box-instance` → `test_box_instance`; models → `test_models`; helpers/common → `test_common`. Sync coverage: `tests/_sync/test_sync_client` + `test_sse_golden`. + +## `exec.session` + +`exec.session()` is a live command session (stdin, PTY, signals, streaming) over +a WebSocket, so it is the one feature not carried by `httpx`. It adds a +`websockets` dependency, imported lazily so it only loads when a session is +opened. + +Both handles are hand-written in `upstash_box/_exec_session.py` rather than +generated: the async handle pumps frames with an asyncio task and the sync +handle with a reader thread, an asymmetry `scripts/generate_sync.py` cannot +produce by token substitution. Frame construction, signal validation, and +decoding are shared between them so the wire protocol has one definition. +`generate_sync.py` maps `AsyncExecSessionHandle`/`open_async_exec_session` to +the sync pair by name. + +Naming follows the SDK's snake_case convention, so the handle is +`end_stdin`/`exec_id` where JS is `endStdin`/`execId`, and callbacks are +`on_stdout`/`on_stderr` taking `bytes`. The sync `wait()` additionally accepts a +`timeout`, since blocking forever on a thread has no async equivalent to +cancellation. + +Note that `scripts/check_parity.py` does **not** gate any of this. Its extractor +walks one level deep — it sees `Box.exec` and `Box.files`, not `exec.session` or +`files.stat` — so nested namespace methods are outside the gate. Treat this file +as the source of truth for namespace-level parity until the extractor recurses. diff --git a/packages/python-sdk/README.md b/packages/python-sdk/README.md index 9586717..e3b0ff3 100644 --- a/packages/python-sdk/README.md +++ b/packages/python-sdk/README.md @@ -162,6 +162,39 @@ print(run.result) # stdout on success, stderr on failure print(run.stdout, run.stderr, run.exit_code) # raw streams + exit code ``` +### Live sessions + +`exec.command` returns after the command finishes. `exec.session` returns as soon +as it starts, so you can write to stdin, resize a PTY, and signal the process +while it runs. + +```python +chunks = [] +session = await box.exec.session( + argv=["sort"], # exact program + args, no shell + on_stdout=chunks.append, # receives bytes as they arrive +) +await session.write("banana\napple\n") +await session.end_stdin() # EOF, so sort finishes +assert await session.wait() == 0 +``` + +Use `cmd="..."` instead of `argv` to go through `bash -lc`, `tty=True` (with +`rows`/`cols`) for a PTY, and `cwd`/`env` to place the process. Control it with +`resize`, `kill(signal)`, `terminate(grace_ms)`, and `close`. + +The session owns the process: closing the handle or losing the connection kills +the command, and sessions cannot be reattached. A context manager makes that +teardown explicit. + +```python +async with await box.exec.session(cmd="npm run dev", tty=True, rows=24, cols=80) as dev: + await dev.write("rs\n") +``` + +The sync client mirrors this without `await`; its `wait(timeout=None)` blocks and +raises `TimeoutError` if the timeout elapses. + ### Files ```python diff --git a/packages/python-sdk/pyproject.toml b/packages/python-sdk/pyproject.toml index 35e0f2a..5a56354 100644 --- a/packages/python-sdk/pyproject.toml +++ b/packages/python-sdk/pyproject.toml @@ -36,6 +36,7 @@ classifiers = [ ] dependencies = [ "httpx>=0.27", + "websockets>=13", "pydantic>=2", "typing-extensions>=4.7", ] diff --git a/packages/python-sdk/scripts/generate_sync.py b/packages/python-sdk/scripts/generate_sync.py index 1a81b48..f6a447d 100644 --- a/packages/python-sdk/scripts/generate_sync.py +++ b/packages/python-sdk/scripts/generate_sync.py @@ -63,6 +63,10 @@ "AsyncSkillsNamespace": "SkillsNamespace", "AsyncLabelsNamespace": "LabelsNamespace", "AsyncClient": "Client", + # Live-exec handles are hand-written in upstash_box/_exec_session.py + # (asyncio task vs reader thread); swap in the sync pair by name. + "AsyncExecSessionHandle": "ExecSessionHandle", + "open_async_exec_session": "open_exec_session", "AsyncIterator": "Iterator", "aiter_bytes": "iter_bytes", "aclose": "close", diff --git a/packages/python-sdk/tests/_async/test_exec_session.py b/packages/python-sdk/tests/_async/test_exec_session.py new file mode 100644 index 0000000..fb23230 --- /dev/null +++ b/packages/python-sdk/tests/_async/test_exec_session.py @@ -0,0 +1,389 @@ +"""Async exec.session handle, driven against a scripted local WebSocket server. + +The protocol helpers are pure and tested directly; everything else runs through +a real socket so the pump task, exit settling, and teardown are exercised rather +than mocked. +""" + +import asyncio +import contextlib +import json +import time + +import pytest +from exec_session_server import replies, start_replies +from websockets.asyncio.server import serve + +from upstash_box import BoxError +from upstash_box._exec_session import ( + build_start_frame, + normalize_signal, + open_async_exec_session, + session_url, +) + +_HANDSHAKE_FAILURES = ( + "__error__", + "__exit__", + "__close__", + "__zero_pid__", + "__no_pid__", +) + + +@pytest.fixture +async def ws_url(): + async def handler(ws): + start = json.loads(await ws.recv()) + for frame in start_replies(start): + await ws.send(json.dumps(frame)) + if start.get("cmd") == "__junk__": + # Unparseable frames, faster than the handshake timeout, then hang + # up so a client that never times out still ends (and fails loudly). + with contextlib.suppress(Exception): + for _ in range(60): + await ws.send("not json") + await asyncio.sleep(0.05) + await ws.close() + return + if start.get("cmd") in _HANDSHAKE_FAILURES: + await ws.close() + return + async for raw in ws: + for frame in replies(json.loads(raw)): + await ws.send(json.dumps(frame)) + + server = await serve(handler, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + try: + yield f"ws://127.0.0.1:{port}" + finally: + server.close() + + +class Collector: + def __init__(self) -> None: + self.out: list[bytes] = [] + self.err: list[bytes] = [] + self._arrived = asyncio.Event() + + def on_stdout(self, data: bytes) -> None: + self.out.append(data) + self._arrived.set() + + def on_stderr(self, data: bytes) -> None: + self.err.append(data) + self._arrived.set() + + async def next_chunk(self) -> None: + await asyncio.wait_for(self._arrived.wait(), 5) + self._arrived.clear() + + +async def open_session(url, *, cmd="run", collector=None, timeout_s=5, **overrides): + fields = { + "argv": None, + "tty": False, + "cwd": "/workspace/home", + "rows": None, + "cols": None, + "env": None, + } + fields.update(overrides) + return await open_async_exec_session( + url=url, + headers={}, + timeout_s=timeout_s, + start=build_start_frame(cmd=cmd, **fields), + on_stdout=collector.on_stdout if collector else None, + on_stderr=collector.on_stderr if collector else None, + ) + + +# ==================== Protocol helpers ==================== + + +def test_start_frame_argv_takes_precedence_over_cmd(): + frame = build_start_frame( + cmd="echo hi", + argv=["/bin/echo", "hi"], + tty=True, + cwd="/workspace/home/sub", + rows=24, + cols=80, + env=["A=1"], + ) + assert frame == { + "type": "start", + "argv": ["/bin/echo", "hi"], + "tty": True, + "cwd": "/workspace/home/sub", + "rows": 24, + "cols": 80, + "env": ["A=1"], + } + assert "cmd" not in frame + + +def test_start_frame_omits_unset_optionals(): + frame = build_start_frame( + cmd="echo hi", argv=None, tty=False, cwd="/workspace/home", rows=None, cols=None, env=None + ) + assert frame == {"type": "start", "cmd": "echo hi", "cwd": "/workspace/home"} + + +def test_start_frame_requires_cmd_or_argv(): + with pytest.raises(BoxError, match="requires cmd or argv"): + build_start_frame( + cmd=None, argv=[], tty=False, cwd="/workspace/home", rows=None, cols=None, env=None + ) + + +@pytest.mark.parametrize( + ("given", "expected"), + [(None, "TERM"), ("int", "INT"), ("SIGKILL", "KILL"), (" sigusr1 ", "USR1")], +) +def test_normalize_signal_accepts_aliases(given, expected): + assert normalize_signal(given) == expected + + +@pytest.mark.parametrize("bad", ["STOP", "SIGSTOP", "9", ""]) +def test_normalize_signal_rejects_unsupported(bad): + with pytest.raises(BoxError, match="unsupported signal"): + normalize_signal(bad) + + +def test_session_url_switches_scheme(): + assert ( + session_url("https://box.upstash.io", "box-1") + == "wss://box.upstash.io/v2/box/box-1/exec-session" + ) + assert ( + session_url("http://localhost:8080", "box-1") + == "ws://localhost:8080/v2/box/box-1/exec-session" + ) + + +# ==================== Handshake ==================== + + +async def test_handshake_exposes_pid_and_exec_id(ws_url): + handle = await open_session(ws_url) + try: + assert handle.pid == 4242 + assert handle.exec_id == "exec-abc" + finally: + await handle.close() + + +async def test_start_frame_reaches_the_server_verbatim(ws_url): + collector = Collector() + handle = await open_session( + ws_url, + cmd=None, + argv=["sleep", "1"], + tty=True, + rows=30, + cols=120, + env=["X=1"], + collector=collector, + ) + try: + await collector.next_chunk() + assert json.loads(collector.out[0]) == { + "type": "start", + "argv": ["sleep", "1"], + "cwd": "/workspace/home", + "env": ["X=1"], + "cols": 120, + "rows": 30, + "tty": True, + } + finally: + await handle.close() + + +@pytest.mark.parametrize("mode", ["__zero_pid__", "__no_pid__"]) +async def test_started_without_a_usable_pid_raises(ws_url, mode): + # A handle whose kill()/terminate() cannot reach the process is worse + # than no handle, so the handshake fails instead. + with pytest.raises(BoxError, match="without a usable pid"): + await open_session(ws_url, cmd=mode) + + +async def test_handshake_deadline_survives_ignored_frames(ws_url): + # The deadline covers the whole handshake, so junk frames cannot extend it. + started_at = time.monotonic() + with pytest.raises(BoxError, match="handshake timed out"): + await open_session(ws_url, cmd="__junk__", timeout_s=0.6) + assert time.monotonic() - started_at < 2.0 + + +async def test_handshake_error_frame_raises(ws_url): + with pytest.raises(BoxError, match="boom"): + await open_session(ws_url, cmd="__error__") + + +async def test_error_frame_after_start_ends_wait_and_hangs_up(ws_url): + handle = await open_session(ws_url, cmd="__late_error__") + assert await handle.wait() == -1 + # Once wait() has settled the caller considers the session over, so the + # client hangs up rather than leaving the process alive behind a live socket. + await asyncio.wait_for(handle._reader, 5) + assert handle._conn.state.name == "CLOSED" + + +async def test_handshake_exit_before_start_raises(ws_url): + with pytest.raises(BoxError, match="exited before start"): + await open_session(ws_url, cmd="__exit__") + + +async def test_handshake_close_before_start_raises(ws_url): + with pytest.raises(BoxError, match="closed before start"): + await open_session(ws_url, cmd="__close__") + + +async def test_connection_failure_raises(ws_url): + with pytest.raises(BoxError, match="connection failed"): + await open_async_exec_session( + url="ws://127.0.0.1:1", + headers={}, + timeout_s=5, + start={"type": "start", "cmd": "x", "cwd": "/"}, + ) + + +# ==================== Live session ==================== + + +async def test_write_reaches_stdin_and_output_is_decoded(ws_url): + collector = Collector() + handle = await open_session(ws_url, collector=collector) + try: + await collector.next_chunk() # start echo + await handle.write("hello ") + await collector.next_chunk() + await handle.write(b"bytes") + await collector.next_chunk() + assert collector.out[1:] == [b"hello ", b"bytes"] + finally: + await handle.close() + + +async def test_end_stdin_drains_stderr_then_exits_zero(ws_url): + collector = Collector() + handle = await open_session(ws_url, collector=collector) + await handle.end_stdin() + assert await handle.wait() == 0 + assert collector.err == [b"eof"] + + +async def test_kill_sends_normalized_signal(ws_url): + collector = Collector() + handle = await open_session(ws_url, collector=collector) + await handle.kill("SIGINT") + assert await handle.wait() == 130 + assert collector.out[-1] == b"sig:INT" + + +async def test_kill_defaults_to_term(ws_url): + collector = Collector() + handle = await open_session(ws_url, collector=collector) + await handle.kill() + await handle.wait() + assert collector.out[-1] == b"sig:TERM" + + +async def test_kill_rejects_unsupported_signal_without_sending(ws_url): + handle = await open_session(ws_url) + try: + with pytest.raises(BoxError, match="unsupported signal"): + await handle.kill("SIGSTOP") + finally: + await handle.close() + + +async def test_terminate_carries_grace_ms(ws_url): + collector = Collector() + handle = await open_session(ws_url, collector=collector) + await handle.terminate(2500) + assert await handle.wait() == 143 + assert collector.out[-1] == b"term:2500" + + +async def test_terminate_omits_non_positive_grace(ws_url): + collector = Collector() + handle = await open_session(ws_url, collector=collector) + await handle.terminate(0) + await handle.wait() + assert collector.out[-1] == b"term:0" + + +async def test_resize_sends_dimensions(ws_url): + collector = Collector() + handle = await open_session(ws_url, tty=True, collector=collector) + try: + await collector.next_chunk() + await handle.resize(40, 100) + await collector.next_chunk() + assert collector.out[-1] == b"size:40x100" + finally: + await handle.close() + + +async def test_writes_after_exit_are_silent_no_ops(ws_url): + collector = Collector() + handle = await open_session(ws_url, collector=collector) + await handle.end_stdin() + assert await handle.wait() == 0 + seen = len(collector.out) + await handle.write("ignored") + await handle.end_stdin() + await handle.resize(1, 1) + await handle.terminate() + await asyncio.sleep(0.05) + assert len(collector.out) == seen + + +async def test_close_settles_wait_as_forced_teardown(ws_url): + handle = await open_session(ws_url) + await handle.close() + assert await handle.wait() == -1 + + +async def test_wait_is_repeatable(ws_url): + handle = await open_session(ws_url) + await handle.end_stdin() + assert await handle.wait() == 0 + assert await handle.wait() == 0 + + +async def test_context_manager_closes_the_session(ws_url): + async with await open_session(ws_url) as handle: + assert handle.pid == 4242 + assert await handle.wait() == -1 + + +async def test_async_callbacks_are_awaited(ws_url): + received: list[bytes] = [] + done = asyncio.Event() + + async def on_stdout(data: bytes) -> None: + await asyncio.sleep(0) + received.append(data) + done.set() + + handle = await open_async_exec_session( + url=ws_url, + headers={}, + timeout_s=5, + start=build_start_frame( + cmd="run", argv=None, tty=False, cwd="/workspace/home", rows=None, cols=None, env=None + ), + on_stdout=on_stdout, + ) + try: + await asyncio.wait_for(done.wait(), 5) + assert json.loads(received[0])["type"] == "start" + finally: + await handle.close() diff --git a/packages/python-sdk/tests/_sync/test_exec_session_sync.py b/packages/python-sdk/tests/_sync/test_exec_session_sync.py new file mode 100644 index 0000000..b8c1f25 --- /dev/null +++ b/packages/python-sdk/tests/_sync/test_exec_session_sync.py @@ -0,0 +1,248 @@ +"""Sync exec.session handle. Hand-written (the sync handle is hand-written too, +not generated), driven against the same scripted server as the async suite so +the two flavors are held to identical wire behavior. +""" + +import contextlib +import json +import threading +import time + +import pytest +from exec_session_server import replies, start_replies +from websockets.sync.server import serve + +from upstash_box import BoxError +from upstash_box._exec_session import build_start_frame, open_exec_session + +_HANDSHAKE_FAILURES = ( + "__error__", + "__exit__", + "__close__", + "__zero_pid__", + "__no_pid__", +) + + +@pytest.fixture +def ws_url(): + def handler(ws): + start = json.loads(ws.recv()) + for frame in start_replies(start): + ws.send(json.dumps(frame)) + if start.get("cmd") == "__junk__": + # See the async suite: junk frames, then hang up. + with contextlib.suppress(Exception): + for _ in range(60): + ws.send("not json") + time.sleep(0.05) + ws.close() + return + if start.get("cmd") in _HANDSHAKE_FAILURES: + ws.close() + return + for raw in ws: + for frame in replies(json.loads(raw)): + ws.send(json.dumps(frame)) + + server = serve(handler, "127.0.0.1", 0) + port = server.socket.getsockname()[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"ws://127.0.0.1:{port}" + finally: + server.shutdown() + thread.join(timeout=5) + + +class Collector: + def __init__(self) -> None: + self.out: list[bytes] = [] + self.err: list[bytes] = [] + self._arrived = threading.Event() + + def on_stdout(self, data: bytes) -> None: + self.out.append(data) + self._arrived.set() + + def on_stderr(self, data: bytes) -> None: + self.err.append(data) + self._arrived.set() + + def next_chunk(self) -> None: + assert self._arrived.wait(5), "no output arrived" + self._arrived.clear() + + +def open_session(url, *, cmd="run", collector=None, timeout_s=5, **overrides): + fields = { + "argv": None, + "tty": False, + "cwd": "/workspace/home", + "rows": None, + "cols": None, + "env": None, + } + fields.update(overrides) + return open_exec_session( + url=url, + headers={}, + timeout_s=timeout_s, + start=build_start_frame(cmd=cmd, **fields), + on_stdout=collector.on_stdout if collector else None, + on_stderr=collector.on_stderr if collector else None, + ) + + +def test_handshake_exposes_pid_and_exec_id(ws_url): + handle = open_session(ws_url) + try: + assert handle.pid == 4242 + assert handle.exec_id == "exec-abc" + finally: + handle.close() + + +def test_start_frame_reaches_the_server_verbatim(ws_url): + collector = Collector() + handle = open_session( + ws_url, cmd=None, argv=["sleep", "1"], tty=True, rows=30, cols=120, collector=collector + ) + try: + collector.next_chunk() + assert json.loads(collector.out[0]) == { + "type": "start", + "argv": ["sleep", "1"], + "cwd": "/workspace/home", + "cols": 120, + "rows": 30, + "tty": True, + } + finally: + handle.close() + + +@pytest.mark.parametrize("mode", ["__zero_pid__", "__no_pid__"]) +def test_started_without_a_usable_pid_raises(ws_url, mode): + # A handle whose kill()/terminate() cannot reach the process is worse + # than no handle, so the handshake fails instead. + with pytest.raises(BoxError, match="without a usable pid"): + open_session(ws_url, cmd=mode) + + +def test_handshake_deadline_survives_ignored_frames(ws_url): + # The deadline covers the whole handshake, so junk frames cannot extend it. + started_at = time.monotonic() + with pytest.raises(BoxError, match="handshake timed out"): + open_session(ws_url, cmd="__junk__", timeout_s=0.6) + assert time.monotonic() - started_at < 2.0 + + +def test_handshake_error_frame_raises(ws_url): + with pytest.raises(BoxError, match="boom"): + open_session(ws_url, cmd="__error__") + + +def test_error_frame_after_start_ends_wait_and_hangs_up(ws_url): + handle = open_session(ws_url, cmd="__late_error__") + assert handle.wait(5) == -1 + # Once wait() has settled the caller considers the session over, so the + # client hangs up rather than leaving the process alive behind a live socket. + handle._reader.join(timeout=5) + assert handle._conn.protocol.state.name == "CLOSED" + + +def test_handshake_close_before_start_raises(ws_url): + with pytest.raises(BoxError, match="closed before start"): + open_session(ws_url, cmd="__close__") + + +def test_connection_failure_raises(): + with pytest.raises(BoxError, match="connection failed"): + open_exec_session( + url="ws://127.0.0.1:1", + headers={}, + timeout_s=5, + start={"type": "start", "cmd": "x", "cwd": "/"}, + ) + + +def test_write_reaches_stdin_and_output_is_decoded(ws_url): + collector = Collector() + handle = open_session(ws_url, collector=collector) + try: + collector.next_chunk() # start echo + handle.write("hello ") + collector.next_chunk() + handle.write(b"bytes") + collector.next_chunk() + assert collector.out[1:] == [b"hello ", b"bytes"] + finally: + handle.close() + + +def test_end_stdin_drains_stderr_then_exits_zero(ws_url): + collector = Collector() + handle = open_session(ws_url, collector=collector) + handle.end_stdin() + assert handle.wait(5) == 0 + assert collector.err == [b"eof"] + + +def test_kill_sends_normalized_signal(ws_url): + collector = Collector() + handle = open_session(ws_url, collector=collector) + handle.kill("SIGINT") + assert handle.wait(5) == 130 + assert collector.out[-1] == b"sig:INT" + + +def test_kill_rejects_unsupported_signal_without_sending(ws_url): + handle = open_session(ws_url) + try: + with pytest.raises(BoxError, match="unsupported signal"): + handle.kill("SIGSTOP") + finally: + handle.close() + + +def test_terminate_carries_grace_ms(ws_url): + collector = Collector() + handle = open_session(ws_url, collector=collector) + handle.terminate(2500) + assert handle.wait(5) == 143 + assert collector.out[-1] == b"term:2500" + + +def test_resize_sends_dimensions(ws_url): + collector = Collector() + handle = open_session(ws_url, tty=True, collector=collector) + try: + collector.next_chunk() + handle.resize(40, 100) + collector.next_chunk() + assert collector.out[-1] == b"size:40x100" + finally: + handle.close() + + +def test_wait_times_out_while_the_process_runs(ws_url): + handle = open_session(ws_url) + try: + with pytest.raises(TimeoutError): + handle.wait(0.2) + finally: + handle.close() + + +def test_close_settles_wait_as_forced_teardown(ws_url): + handle = open_session(ws_url) + handle.close() + assert handle.wait(5) == -1 + + +def test_context_manager_closes_the_session(ws_url): + with open_session(ws_url) as handle: + assert handle.pid == 4242 + assert handle.wait(5) == -1 diff --git a/packages/python-sdk/tests/exec_session_server.py b/packages/python-sdk/tests/exec_session_server.py new file mode 100644 index 0000000..bbe0a0c --- /dev/null +++ b/packages/python-sdk/tests/exec_session_server.py @@ -0,0 +1,61 @@ +"""Scripted exec-session server backing the async and sync handle tests. + +The reply table is shared so both flavors are driven by identical server +behavior; only the socket plumbing differs per test module. A sentinel ``cmd`` +in the start frame selects a failure mode for the handshake tests. +""" + +import base64 +import json +from typing import Any, Dict, List + +STARTED = {"type": "started", "pid": 4242, "execId": "exec-abc"} + + +def b64(text: str) -> str: + return base64.b64encode(text.encode("utf-8")).decode("ascii") + + +def start_replies(start: Dict[str, Any]) -> List[Dict[str, Any]]: + """Frames sent in response to the opening frame. Echoes the start frame back + as stdout so tests can assert exactly what the client put on the wire.""" + cmd = start.get("cmd") + if cmd == "__error__": + return [{"type": "error", "message": "boom"}] + if cmd == "__exit__": + return [{"type": "exit", "code": 7}] + if cmd in ("__close__", "__junk__"): + # __junk__ sends nothing here; the fixture then streams unparseable + # frames so a per-receive handshake timeout would never fire. + return [] + if cmd == "__late_error__": + # Starts cleanly, then fails. The server stays connected on purpose, so + # a test can prove the client is the one that hangs up. + return [STARTED, {"type": "error", "message": "late boom"}] + if cmd == "__zero_pid__": + return [{"type": "started", "pid": 0, "execId": "exec-abc"}] + if cmd == "__no_pid__": + return [{"type": "started", "execId": "exec-abc"}] + return [STARTED, {"type": "stdout", "data": b64(json.dumps(start, sort_keys=True))}] + + +def replies(msg: Dict[str, Any]) -> List[Dict[str, Any]]: + """Frames sent in response to one post-handshake client frame.""" + kind = msg.get("type") + if kind == "stdin": + return [{"type": "stdout", "data": msg["data"]}] + if kind == "stdin_close": + return [{"type": "stderr", "data": b64("eof")}, {"type": "exit", "code": 0}] + if kind == "signal": + return [ + {"type": "stdout", "data": b64(f"sig:{msg['signal']}")}, + {"type": "exit", "code": 130}, + ] + if kind == "terminate": + return [ + {"type": "stdout", "data": b64(f"term:{msg.get('graceMs', 0)}")}, + {"type": "exit", "code": 143}, + ] + if kind == "resize": + return [{"type": "stdout", "data": b64(f"size:{msg['rows']}x{msg['cols']}")}] + return [] diff --git a/packages/python-sdk/tests/integration/conftest.py b/packages/python-sdk/tests/integration/conftest.py index d797466..d1c5cdc 100644 --- a/packages/python-sdk/tests/integration/conftest.py +++ b/packages/python-sdk/tests/integration/conftest.py @@ -27,9 +27,20 @@ def pytest_collection_modifyitems(config, items): item.add_marker(skip) -@pytest.fixture -def opts() -> dict: +def _opts() -> dict: kwargs = {"api_key": API_KEY} if BASE_URL: kwargs["base_url"] = BASE_URL return kwargs + + +@pytest.fixture +def opts() -> dict: + return _opts() + + +@pytest.fixture(scope="module") +def module_opts() -> dict: + """Same credentials, module-scoped, for suites that share one box across + tests instead of creating one per test.""" + return _opts() diff --git a/packages/python-sdk/tests/integration/test_exec_session_async.py b/packages/python-sdk/tests/integration/test_exec_session_async.py new file mode 100644 index 0000000..d51eba3 --- /dev/null +++ b/packages/python-sdk/tests/integration/test_exec_session_async.py @@ -0,0 +1,230 @@ +"""Live exec sessions against a real box. Mirrors the `@upstash/box` integration +suite so both SDKs are held to the same server behavior.""" + +import asyncio + +import pytest +import pytest_asyncio + +from upstash_box import AsyncBox, BoxError + +# One box for the whole module, so the loop must outlive each test. +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="module")] + + +class Sink: + """Accumulates decoded output from a session's callbacks.""" + + def __init__(self) -> None: + self.out = "" + self.err = "" + + def on_stdout(self, data: bytes) -> None: + self.out += data.decode("utf-8", "replace") + + def on_stderr(self, data: bytes) -> None: + self.err += data.decode("utf-8", "replace") + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def box(module_opts): + created = await AsyncBox.create(**module_opts) + try: + yield created + finally: + try: + await created.delete() + except Exception: + pass + + +async def test_streams_stdout_and_stderr_separately_with_exit_code(box): + sink = Sink() + session = await box.exec.session( + argv=["sh", "-c", "echo to-stdout; echo to-stderr 1>&2; exit 42"], + on_stdout=sink.on_stdout, + on_stderr=sink.on_stderr, + ) + assert session.pid > 0 + assert session.exec_id != "" + assert await session.wait() == 42 + assert sink.out.strip() == "to-stdout" + assert sink.err.strip() == "to-stderr" + + +async def test_argv_runs_without_a_shell(box): + sink = Sink() + session = await box.exec.session(argv=["echo", "$HOME; rm -rf /"], on_stdout=sink.on_stdout) + await session.wait() + # A shell would expand $HOME and treat `;` as a separator. + assert sink.out.strip() == "$HOME; rm -rf /" + + +async def test_cmd_runs_through_a_shell(box): + sink = Sink() + session = await box.exec.session(cmd="echo shell-$((1+1))", on_stdout=sink.on_stdout) + await session.wait() + assert sink.out.strip() == "shell-2" + + +async def test_stdin_and_end_stdin_finish_an_eof_reading_command(box): + sink = Sink() + session = await box.exec.session(argv=["sort"], on_stdout=sink.on_stdout) + await session.write("banana\napple\ncherry\n") + await session.end_stdin() + + assert await session.wait() == 0 + assert sink.out == "apple\nbanana\ncherry\n" + + +async def test_honors_cwd_and_overlays_env(box): + await box.files.mkdir("session-proj/src", parents=True) + sink = Sink() + session = await box.exec.session( + argv=["sh", "-c", "pwd; echo $MY_VAR"], + cwd="session-proj/src", + env=["MY_VAR=from-test"], + on_stdout=sink.on_stdout, + ) + await session.wait() + + assert "/workspace/home/session-proj/src" in sink.out + assert "from-test" in sink.out + await box.files.remove("session-proj", recursive=True) + + +async def test_drops_blocked_env_keys_but_passes_ordinary_ones(box): + sink = Sink() + session = await box.exec.session( + argv=["sh", "-c", "echo LD=[$LD_PRELOAD] SAFE=[$SAFE]"], + env=["LD_PRELOAD=/tmp/evil.so", "SAFE=yes"], + on_stdout=sink.on_stdout, + ) + await session.wait() + + assert "LD=[]" in sink.out + assert "SAFE=[yes]" in sink.out + + +async def test_terminate_stops_a_long_running_process(box): + session = await box.exec.session(argv=["sleep", "300"]) + await session.terminate(1000) + assert await session.wait() != 0 + + +async def test_kill_reaps_the_whole_process_tree(box): + async def running() -> str: + run = await box.exec.command( + 'c=0; for d in /proc/[0-9]*; do [ "$(cat "$d/comm" 2>/dev/null)" = "sleep" ] ' + '&& grep -qs 4711 "$d/cmdline" && c=$((c+1)); done; echo $c' + ) + return run.result.strip() + + session = await box.exec.session(cmd="sleep 4711 & sleep 4712 & wait") + await asyncio.sleep(0.8) + assert await running() != "0" + + await session.kill("TERM") + await session.wait() + await asyncio.sleep(0.5) + assert await running() == "0" + + +async def test_allocates_a_real_pty_at_the_requested_size(box): + sink = Sink() + session = await box.exec.session( + tty=True, + rows=24, + cols=80, + cmd="tty; stty size; read line; echo GOT=$line; exit 0", + on_stdout=sink.on_stdout, + ) + await asyncio.sleep(0.4) + await session.write("hello-pty\n") + + assert await session.wait() == 0 + assert "/dev/pts/" in sink.out + # Size must be right from the first read, not applied after the process starts. + assert "24 80" in sink.out + assert "GOT=hello-pty" in sink.out + + +async def test_keeps_a_long_lived_process_for_multiple_round_trips(box): + sink = Sink() + session = await box.exec.session(argv=["cat"], on_stdout=sink.on_stdout) + for msg in ("req-1\n", "req-2\n", "req-3\n"): + await session.write(msg) + await asyncio.sleep(0.15) + assert "req-1" in sink.out + assert "req-2" in sink.out + assert "req-3" in sink.out + + await session.kill("KILL") + await session.wait() + + +async def test_runs_sessions_concurrently_without_crosstalk(box): + async def worker(n: int): + sink = Sink() + session = await box.exec.session( + argv=["sh", "-c", f"sleep 0.{n}; echo worker-{n}"], on_stdout=sink.on_stdout + ) + return n, await session.wait(), sink.out.strip() + + for n, code, out in await asyncio.gather(*(worker(n) for n in (1, 2, 3, 4))): + assert code == 0 + assert out == f"worker-{n}" + + +async def test_does_not_leak_env_between_sessions(box): + first = await box.exec.session(argv=["sh", "-c", "export LEAK=nope; true"]) + await first.wait() + + sink = Sink() + session = await box.exec.session( + argv=["sh", "-c", "echo LEAK=[$LEAK]"], on_stdout=sink.on_stdout + ) + await session.wait() + assert "LEAK=[]" in sink.out + + +async def test_close_ends_the_session_and_stops_the_process(box): + session = await box.exec.session(argv=["sleep", "600"]) + pid = session.pid + await session.close() + await session.wait() + + await asyncio.sleep(1.5) + alive = await box.exec.command(f"[ -d /proc/{pid} ] && echo yes || echo no") + assert alive.result.strip() == "no" + + +async def test_rejects_empty_command_locally_and_unsupported_signal(box): + with pytest.raises(BoxError, match="requires cmd or argv"): + await box.exec.session() + + session = await box.exec.session(argv=["sleep", "60"]) + with pytest.raises(BoxError, match="unsupported signal"): + await session.kill("BOGUS") + await session.close() + await session.wait() + + +async def test_session_writes_are_visible_to_the_files_api(box): + session = await box.exec.session(argv=["sh", "-c", "echo written-by-session > session-out.txt"]) + assert await session.wait() == 0 + + assert (await box.files.read("session-out.txt")).strip() == "written-by-session" + assert (await box.files.stat("session-out.txt")).type == "file" + await box.files.remove("session-out.txt") + + +async def test_context_manager_tears_the_session_down(box): + async with await box.exec.session(argv=["sleep", "600"]) as session: + pid = session.pid + assert pid > 0 + await session.wait() + + await asyncio.sleep(1.5) + alive = await box.exec.command(f"[ -d /proc/{pid} ] && echo yes || echo no") + assert alive.result.strip() == "no" diff --git a/packages/python-sdk/tests/integration/test_sync_subset.py b/packages/python-sdk/tests/integration/test_sync_subset.py index 9044334..3d0d0bd 100644 --- a/packages/python-sdk/tests/integration/test_sync_subset.py +++ b/packages/python-sdk/tests/integration/test_sync_subset.py @@ -33,5 +33,14 @@ def test_sync_create_exec_agent_stream_delete(opts, tmp_path): chunks += 1 assert chunks > 0 assert stream.status == "completed" + + # Live session on the sync handle: reader thread, stdin, EOF, exit code. + chunks_out = [] + session = box.exec.session(argv=["sort"], on_stdout=chunks_out.append) + assert session.pid > 0 + session.write("pear\napple\n") + session.end_stdin() + assert session.wait(30) == 0 + assert b"".join(chunks_out).decode() == "apple\npear\n" finally: box.delete() diff --git a/packages/python-sdk/upstash_box/__init__.py b/packages/python-sdk/upstash_box/__init__.py index 8185336..29e518e 100644 --- a/packages/python-sdk/upstash_box/__init__.py +++ b/packages/python-sdk/upstash_box/__init__.py @@ -22,6 +22,7 @@ AsyncTab, ) from ._common import infer_default_provider +from ._exec_session import AsyncExecSessionHandle, ExecSessionHandle from ._sync import ( AgentNamespace, Box, @@ -126,6 +127,8 @@ ) __all__ = [ + "AsyncExecSessionHandle", + "ExecSessionHandle", "__version__", # Clients (sync canonical + async variants) "Box", diff --git a/packages/python-sdk/upstash_box/_async/client.py b/packages/python-sdk/upstash_box/_async/client.py index cb1375c..f211b6f 100644 --- a/packages/python-sdk/upstash_box/_async/client.py +++ b/packages/python-sdk/upstash_box/_async/client.py @@ -33,6 +33,13 @@ from typing_extensions import Unpack from .. import _common as common +from .._exec_session import ( + AsyncExecSessionHandle, + StdoutCallback, + build_start_frame, + open_async_exec_session, + session_url, +) from ..errors import BoxError from ..types import ( Agent, @@ -296,6 +303,46 @@ async def stream_code( ) -> AsyncStreamRun[str]: return await self._box._exec_stream_code(code, lang, timeout) + async def session( + self, + *, + cmd: Optional[str] = None, + argv: Optional[List[str]] = None, + tty: bool = False, + cwd: Optional[str] = None, + rows: Optional[int] = None, + cols: Optional[int] = None, + env: Optional[List[str]] = None, + on_stdout: Optional[StdoutCallback] = None, + on_stderr: Optional[StdoutCallback] = None, + ) -> AsyncExecSessionHandle: + """Start a live command session and return once the process is running. + + Unlike ``command``, which returns after the command finishes, a session + stays open: write to stdin, resize a PTY, signal the process tree, and + receive output as it is produced. + + ``argv`` is the exact program and arguments, run without a shell, and + takes precedence over ``cmd``, which is run via ``bash -lc``. ``tty`` + allocates a PTY, which merges stderr into stdout. ``env`` entries are + ``KEY=VALUE`` strings overlaid on the box environment. ``on_stdout`` and + ``on_stderr`` receive raw ``bytes`` as they arrive. + + The returned handle owns the process: closing it, or losing the + connection, kills the command. + """ + return await self._box._exec_session( + cmd=cmd, + argv=argv, + tty=tty, + cwd=cwd, + rows=rows, + cols=cols, + env=env, + on_stdout=on_stdout, + on_stderr=on_stderr, + ) + class AsyncFilesNamespace: def __init__(self, box: "AsyncBox") -> None: @@ -1202,6 +1249,38 @@ async def _exec_stream_code(self, code, lang, timeout) -> AsyncStreamRun[str]: body["folder"] = folder return await self._exec_stream_request(f"/v2/box/{self.id}/code-stream", body, "code") + async def _exec_session( + self, + *, + cmd: Optional[str] = None, + argv: Optional[List[str]] = None, + tty: bool = False, + cwd: Optional[str] = None, + rows: Optional[int] = None, + cols: Optional[int] = None, + env: Optional[List[str]] = None, + on_stdout: Optional[StdoutCallback] = None, + on_stderr: Optional[StdoutCallback] = None, + ) -> AsyncExecSessionHandle: + start = build_start_frame( + cmd=cmd, + argv=argv, + tty=tty, + # Default to the box's current directory (honoring cd()), matching exec.command. + cwd=self._resolve_path(cwd) if cwd else self._cwd, + rows=rows, + cols=cols, + env=env, + ) + return await open_async_exec_session( + url=session_url(self._base_url, self.id), + headers=dict(self._headers), + timeout_s=_ms_to_seconds(self._timeout_ms), + start=start, + on_stdout=on_stdout, + on_stderr=on_stderr, + ) + async def _exec_stream_request(self, path, body, type_) -> AsyncStreamRun[str]: start = time.time() * 1000 run: AsyncStreamRun[str] = AsyncStreamRun(self, type_) diff --git a/packages/python-sdk/upstash_box/_exec_session.py b/packages/python-sdk/upstash_box/_exec_session.py new file mode 100644 index 0000000..093956b --- /dev/null +++ b/packages/python-sdk/upstash_box/_exec_session.py @@ -0,0 +1,513 @@ +"""Live exec sessions over WebSocket. + +Both flavors are hand-written. The async handle pumps frames with an asyncio +task and the sync handle with a reader thread, an asymmetry that +scripts/generate_sync.py cannot produce by token substitution. Everything that +defines the wire protocol (frame construction, validation, decoding) is shared +below so the two handles cannot drift apart. +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import inspect +import json +import re +import threading +import time +from typing import Any, Callable, Dict, List, Optional, Union + +from .errors import BoxError + +StdoutCallback = Callable[[bytes], Any] + +_SIGNALS = frozenset({"TERM", "KILL", "INT", "HUP", "TSTP", "QUIT", "USR1", "USR2"}) + + +def session_url(base_url: str, box_id: str) -> str: + return re.sub(r"^http", "ws", base_url) + f"/v2/box/{box_id}/exec-session" + + +def build_start_frame( + *, + cmd: Optional[str], + argv: Optional[List[str]], + tty: bool, + cwd: str, + rows: Optional[int], + cols: Optional[int], + env: Optional[List[str]], +) -> Dict[str, Any]: + """Assemble the opening frame. ``argv`` takes precedence over ``cmd``.""" + has_argv = bool(argv) + if not has_argv and not cmd: + raise BoxError("exec.session requires cmd or argv") + start: Dict[str, Any] = {"type": "start"} + if has_argv: + start["argv"] = list(argv or []) + else: + start["cmd"] = cmd + if tty: + start["tty"] = True + start["cwd"] = cwd + if rows: + start["rows"] = rows + if cols: + start["cols"] = cols + if env: + start["env"] = list(env) + return start + + +def normalize_signal(signal: Optional[str]) -> str: + sig = ("TERM" if signal is None else signal).strip().upper() + if sig.startswith("SIG"): + sig = sig[3:] + if sig not in _SIGNALS: + raise BoxError(f"unsupported signal: {signal}") + return sig + + +def stdin_frame(data: Union[str, bytes]) -> Dict[str, Any]: + payload = data.encode("utf-8") if isinstance(data, str) else bytes(data) + return {"type": "stdin", "data": base64.b64encode(payload).decode("ascii")} + + +def terminate_frame(grace_ms: Optional[int]) -> Dict[str, Any]: + frame: Dict[str, Any] = {"type": "terminate"} + if grace_ms and grace_ms > 0: + frame["graceMs"] = grace_ms + return frame + + +def parse_frame(raw: Union[str, bytes]) -> Optional[Dict[str, Any]]: + """Decode one server frame, ignoring anything unparseable.""" + try: + frame = json.loads(raw if isinstance(raw, str) else raw.decode("utf-8")) + except Exception: + return None + return frame if isinstance(frame, dict) else None + + +def frame_payload(frame: Dict[str, Any]) -> bytes: + data = frame.get("data") + if not isinstance(data, str): + return b"" + try: + return base64.b64decode(data) + except Exception: + return b"" + + +def exit_code_of(frame: Dict[str, Any]) -> int: + code = frame.get("code") + return code if isinstance(code, int) else -1 + + +def _handshake_error(frame: Dict[str, Any]) -> BoxError: + return BoxError(f"exec-session error: {frame.get('message')}") + + +def _started_fields(frame: Dict[str, Any]) -> "tuple[int, str]": + """Read the started frame. A session whose process cannot be signaled is + useless, so a missing or non-positive pid fails the handshake rather than + producing a handle whose ``kill``/``terminate`` would go nowhere.""" + pid = frame.get("pid") + if not isinstance(pid, int) or isinstance(pid, bool) or pid <= 0: + raise BoxError("exec-session started without a usable pid") + exec_id = frame.get("execId") + return (pid, exec_id if isinstance(exec_id, str) else "") + + +# ==================== Async ==================== + + +class AsyncExecSessionHandle: + """A live command session. + + ``session()`` returns this once the process has started. Output flows to the + ``on_stdout``/``on_stderr`` callbacks passed to ``session()``. + + The session owns the process: losing the connection kills it. ``close()``, a + dropped network link, or exiting the program all terminate the command + rather than leaving it running in the box, and sessions cannot be + reattached. Use ``wait()`` to run something to completion. + + Callbacks run on the task pumping the socket, so a slow callback delays + later output. An ``async`` callback is awaited, so it can do I/O without + blocking the loop. + """ + + pid: int + """In-box (container-namespace) PID, always non-zero: a session whose + process cannot be signaled fails the handshake instead.""" + + exec_id: str + """Server-side exec id.""" + + def __init__( + self, + conn: Any, + pid: int, + exec_id: str, + on_stdout: Optional[StdoutCallback], + on_stderr: Optional[StdoutCallback], + ) -> None: + self.pid = pid + self.exec_id = exec_id + self._conn = conn + self._on_stdout = on_stdout + self._on_stderr = on_stderr + self._exit: asyncio.Future[int] = asyncio.get_running_loop().create_future() + self._reader = asyncio.ensure_future(self._pump()) + + async def _dispatch(self, cb: Optional[StdoutCallback], frame: Dict[str, Any]) -> None: + if cb is None: + return + result = cb(frame_payload(frame)) + if inspect.isawaitable(result): + await result + + async def _pump(self) -> None: + try: + async for raw in self._conn: + frame = parse_frame(raw) + if frame is None: + continue + kind = frame.get("type") + if kind == "stdout": + await self._dispatch(self._on_stdout, frame) + elif kind == "stderr": + await self._dispatch(self._on_stderr, frame) + elif kind == "exit": + self._settle(exit_code_of(frame)) + break + elif kind == "error": + self._settle(-1) + break + except Exception: + pass + finally: + self._settle(-1) + with contextlib.suppress(Exception): + await self._conn.close() + + def _settle(self, code: int) -> None: + if not self._exit.done(): + self._exit.set_result(code) + + async def _send(self, frame: Dict[str, Any]) -> None: + try: + await self._conn.send(json.dumps(frame)) + except Exception as exc: + raise BoxError(f"exec-session send failed: {exc}") from exc + + async def write(self, data: Union[str, bytes]) -> None: + """Write bytes to the process stdin.""" + if self._exit.done(): + return + await self._send(stdin_frame(data)) + + async def end_stdin(self) -> None: + """Close stdin (send EOF). A command that reads until EOF (``cat``, + ``sort``) then exits on its own; output keeps flowing until it does.""" + if self._exit.done(): + return + await self._send({"type": "stdin_close"}) + + async def resize(self, rows: int, cols: int) -> None: + """Resize the PTY (TTY sessions).""" + if self._exit.done(): + return + await self._send({"type": "resize", "rows": rows, "cols": cols}) + + async def kill(self, signal: Optional[str] = None) -> None: + """Send a signal to the process tree. Defaults to ``TERM``.""" + if self._exit.done(): + return + await self._send({"type": "signal", "signal": normalize_signal(signal)}) + + async def terminate(self, grace_ms: Optional[int] = None) -> None: + """Graceful stop driven server-side: SIGTERM now, then SIGKILL after + ``grace_ms`` (default is the server's grace) if still running. + + Only the first call starts the sequence; later ones are ignored, so the + grace cannot be changed once it is running. Use ``kill("KILL")`` to stop + the process immediately instead.""" + if self._exit.done(): + return + await self._send(terminate_frame(grace_ms)) + + async def wait(self) -> int: + """Wait for the process to finish and return its exit code (``-1`` if it + was still running at a forced teardown).""" + return await asyncio.shield(self._exit) + + async def close(self) -> None: + """Close the connection, terminating the process if still running.""" + with contextlib.suppress(Exception): + await self._conn.close() + if self._reader is not asyncio.current_task(): + with contextlib.suppress(Exception): + await self._reader + self._settle(-1) + + async def __aenter__(self) -> "AsyncExecSessionHandle": + return self + + async def __aexit__(self, *_exc: Any) -> None: + await self.close() + + +async def open_async_exec_session( + *, + url: str, + headers: Dict[str, str], + timeout_s: Optional[float], + start: Dict[str, Any], + on_stdout: Optional[StdoutCallback] = None, + on_stderr: Optional[StdoutCallback] = None, +) -> AsyncExecSessionHandle: + try: + from websockets.asyncio.client import connect + except ImportError as exc: # pragma: no cover + raise BoxError("exec.session requires the 'websockets' package") from exc + + try: + conn = await connect(url, additional_headers=headers, open_timeout=timeout_s, max_size=None) + except Exception as exc: + raise BoxError(f"exec-session connection failed: {exc}") from exc + + handshake_failed = True + try: + await conn.send(json.dumps(start)) + # One deadline for the whole handshake. Frames that are not "started" + # are skipped below, so a per-receive timeout would restart on every + # ignored frame and a chatty peer could keep session() pending forever. + loop = asyncio.get_running_loop() + deadline = None if timeout_s is None else loop.time() + timeout_s + while True: + try: + remaining = None if deadline is None else max(0.0, deadline - loop.time()) + raw = await asyncio.wait_for(conn.recv(), remaining) + except asyncio.TimeoutError as exc: + raise BoxError("exec.session handshake timed out") from exc + except Exception as exc: + raise BoxError("exec-session closed before start") from exc + frame = parse_frame(raw) + if frame is None: + continue + kind = frame.get("type") + if kind == "started": + pid, exec_id = _started_fields(frame) + handle = AsyncExecSessionHandle(conn, pid, exec_id, on_stdout, on_stderr) + handshake_failed = False + return handle + if kind == "error": + raise _handshake_error(frame) + if kind == "exit": + raise BoxError("exec-session exited before start") + finally: + if handshake_failed: + with contextlib.suppress(Exception): + await conn.close() + + +# ==================== Sync ==================== + + +class ExecSessionHandle: + """A live command session. + + ``session()`` returns this once the process has started. Output flows to the + ``on_stdout``/``on_stderr`` callbacks passed to ``session()``, invoked on a + background reader thread. + + The session owns the process: losing the connection kills it. ``close()``, a + dropped network link, or exiting the program all terminate the command + rather than leaving it running in the box, and sessions cannot be + reattached. Use ``wait()`` to run something to completion. + + Callbacks run on the background reader thread, which is also what delivers + the exit frame. Keep them short: blocking there stalls the stream, and + calling ``wait()`` from inside one deadlocks, because the exit it waits for + can only arrive on the thread it is blocking. Hand work to your own queue + instead. + """ + + pid: int + """In-box (container-namespace) PID, always non-zero: a session whose + process cannot be signaled fails the handshake instead.""" + + exec_id: str + """Server-side exec id.""" + + def __init__( + self, + conn: Any, + pid: int, + exec_id: str, + on_stdout: Optional[StdoutCallback], + on_stderr: Optional[StdoutCallback], + ) -> None: + self.pid = pid + self.exec_id = exec_id + self._conn = conn + self._on_stdout = on_stdout + self._on_stderr = on_stderr + self._exit_code = -1 + self._exited = threading.Event() + self._send_lock = threading.Lock() + self._reader = threading.Thread(target=self._pump, daemon=True) + self._reader.start() + + def _dispatch(self, cb: Optional[StdoutCallback], frame: Dict[str, Any]) -> None: + if cb is not None: + cb(frame_payload(frame)) + + def _pump(self) -> None: + try: + for raw in self._conn: + frame = parse_frame(raw) + if frame is None: + continue + kind = frame.get("type") + if kind == "stdout": + self._dispatch(self._on_stdout, frame) + elif kind == "stderr": + self._dispatch(self._on_stderr, frame) + elif kind == "exit": + self._settle(exit_code_of(frame)) + break + elif kind == "error": + self._settle(-1) + break + except Exception: + pass + finally: + self._settle(-1) + with contextlib.suppress(Exception): + self._conn.close() + + def _settle(self, code: int) -> None: + if not self._exited.is_set(): + self._exit_code = code + self._exited.set() + + def _send(self, frame: Dict[str, Any]) -> None: + try: + with self._send_lock: + self._conn.send(json.dumps(frame)) + except Exception as exc: + raise BoxError(f"exec-session send failed: {exc}") from exc + + def write(self, data: Union[str, bytes]) -> None: + """Write bytes to the process stdin.""" + if self._exited.is_set(): + return + self._send(stdin_frame(data)) + + def end_stdin(self) -> None: + """Close stdin (send EOF). A command that reads until EOF (``cat``, + ``sort``) then exits on its own; output keeps flowing until it does.""" + if self._exited.is_set(): + return + self._send({"type": "stdin_close"}) + + def resize(self, rows: int, cols: int) -> None: + """Resize the PTY (TTY sessions).""" + if self._exited.is_set(): + return + self._send({"type": "resize", "rows": rows, "cols": cols}) + + def kill(self, signal: Optional[str] = None) -> None: + """Send a signal to the process tree. Defaults to ``TERM``.""" + if self._exited.is_set(): + return + self._send({"type": "signal", "signal": normalize_signal(signal)}) + + def terminate(self, grace_ms: Optional[int] = None) -> None: + """Graceful stop driven server-side: SIGTERM now, then SIGKILL after + ``grace_ms`` (default is the server's grace) if still running. + + Only the first call starts the sequence; later ones are ignored, so the + grace cannot be changed once it is running. Use ``kill("KILL")`` to stop + the process immediately instead.""" + if self._exited.is_set(): + return + self._send(terminate_frame(grace_ms)) + + def wait(self, timeout: Optional[float] = None) -> int: + """Wait for the process to finish and return its exit code (``-1`` if it + was still running at a forced teardown). Raises ``TimeoutError`` if + ``timeout`` elapses first.""" + if not self._exited.wait(timeout): + raise TimeoutError("exec.session wait timed out") + return self._exit_code + + def close(self) -> None: + """Close the connection, terminating the process if still running.""" + with contextlib.suppress(Exception): + self._conn.close() + if self._reader is not threading.current_thread(): + self._reader.join(timeout=5) + self._settle(-1) + + def __enter__(self) -> "ExecSessionHandle": + return self + + def __exit__(self, *_exc: Any) -> None: + self.close() + + +def open_exec_session( + *, + url: str, + headers: Dict[str, str], + timeout_s: Optional[float], + start: Dict[str, Any], + on_stdout: Optional[StdoutCallback] = None, + on_stderr: Optional[StdoutCallback] = None, +) -> ExecSessionHandle: + try: + from websockets.sync.client import connect + except ImportError as exc: # pragma: no cover + raise BoxError("exec.session requires the 'websockets' package") from exc + + try: + conn = connect(url, additional_headers=headers, open_timeout=timeout_s, max_size=None) + except Exception as exc: + raise BoxError(f"exec-session connection failed: {exc}") from exc + + handshake_failed = True + try: + conn.send(json.dumps(start)) + # One deadline for the whole handshake; see the async note above. + deadline = None if timeout_s is None else time.monotonic() + timeout_s + while True: + try: + remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) + raw = conn.recv(timeout=remaining) + except TimeoutError as exc: + raise BoxError("exec.session handshake timed out") from exc + except Exception as exc: + raise BoxError("exec-session closed before start") from exc + frame = parse_frame(raw) + if frame is None: + continue + kind = frame.get("type") + if kind == "started": + pid, exec_id = _started_fields(frame) + handle = ExecSessionHandle(conn, pid, exec_id, on_stdout, on_stderr) + handshake_failed = False + return handle + if kind == "error": + raise _handshake_error(frame) + if kind == "exit": + raise BoxError("exec-session exited before start") + finally: + if handshake_failed: + with contextlib.suppress(Exception): + conn.close() diff --git a/packages/python-sdk/upstash_box/_sync/client.py b/packages/python-sdk/upstash_box/_sync/client.py index 7d104e1..43f9bc2 100644 --- a/packages/python-sdk/upstash_box/_sync/client.py +++ b/packages/python-sdk/upstash_box/_sync/client.py @@ -32,6 +32,13 @@ from typing_extensions import Unpack from .. import _common as common +from .._exec_session import ( + ExecSessionHandle, + StdoutCallback, + build_start_frame, + open_exec_session, + session_url, +) from ..errors import BoxError from ..types import ( Agent, @@ -291,6 +298,46 @@ def stream_code( ) -> StreamRun[str]: return self._box._exec_stream_code(code, lang, timeout) + def session( + self, + *, + cmd: Optional[str] = None, + argv: Optional[List[str]] = None, + tty: bool = False, + cwd: Optional[str] = None, + rows: Optional[int] = None, + cols: Optional[int] = None, + env: Optional[List[str]] = None, + on_stdout: Optional[StdoutCallback] = None, + on_stderr: Optional[StdoutCallback] = None, + ) -> ExecSessionHandle: + """Start a live command session and return once the process is running. + + Unlike ``command``, which returns after the command finishes, a session + stays open: write to stdin, resize a PTY, signal the process tree, and + receive output as it is produced. + + ``argv`` is the exact program and arguments, run without a shell, and + takes precedence over ``cmd``, which is run via ``bash -lc``. ``tty`` + allocates a PTY, which merges stderr into stdout. ``env`` entries are + ``KEY=VALUE`` strings overlaid on the box environment. ``on_stdout`` and + ``on_stderr`` receive raw ``bytes`` as they arrive. + + The returned handle owns the process: closing it, or losing the + connection, kills the command. + """ + return self._box._exec_session( + cmd=cmd, + argv=argv, + tty=tty, + cwd=cwd, + rows=rows, + cols=cols, + env=env, + on_stdout=on_stdout, + on_stderr=on_stderr, + ) + class FilesNamespace: def __init__(self, box: "Box") -> None: @@ -1191,6 +1238,38 @@ def _exec_stream_code(self, code, lang, timeout) -> StreamRun[str]: body["folder"] = folder return self._exec_stream_request(f"/v2/box/{self.id}/code-stream", body, "code") + def _exec_session( + self, + *, + cmd: Optional[str] = None, + argv: Optional[List[str]] = None, + tty: bool = False, + cwd: Optional[str] = None, + rows: Optional[int] = None, + cols: Optional[int] = None, + env: Optional[List[str]] = None, + on_stdout: Optional[StdoutCallback] = None, + on_stderr: Optional[StdoutCallback] = None, + ) -> ExecSessionHandle: + start = build_start_frame( + cmd=cmd, + argv=argv, + tty=tty, + # Default to the box's current directory (honoring cd()), matching exec.command. + cwd=self._resolve_path(cwd) if cwd else self._cwd, + rows=rows, + cols=cols, + env=env, + ) + return open_exec_session( + url=session_url(self._base_url, self.id), + headers=dict(self._headers), + timeout_s=_ms_to_seconds(self._timeout_ms), + start=start, + on_stdout=on_stdout, + on_stderr=on_stderr, + ) + def _exec_stream_request(self, path, body, type_) -> StreamRun[str]: start = time.time() * 1000 run: StreamRun[str] = StreamRun(self, type_) diff --git a/packages/sdk/README.md b/packages/sdk/README.md index f1fd598..dd3f8a0 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -154,6 +154,44 @@ const run = await box.exec.command("node index.js"); console.log(run.result); ``` +#### `box.exec.session(options): Promise` + +Start a live command session. `exec.command` resolves once the command has +finished; `session` resolves as soon as it starts, so you can write to stdin, +resize a PTY, and signal the process while it runs. + +This one is Node-only. Authentication travels in a request header, and browsers +cannot set headers on a WebSocket handshake. + +```ts +let out = ""; +const session = await box.exec.session({ + argv: ["sort"], // exact program + args, no shell + onStdout: (b) => (out += Buffer.from(b).toString()), +}); +session.write("banana\napple\n"); +session.endStdin(); // EOF, so sort finishes +await session.wait(); // 0 +``` + +Pass `cmd` instead of `argv` to run through `bash -lc`, `tty: true` (with +`rows` / `cols`) to allocate a PTY sized correctly from the first read, and +`cwd` / `env` to place the process. `env` entries are `KEY=VALUE` strings +overlaid on the box environment. + +The handle exposes `pid` and `execId`, plus `write`, `endStdin`, `resize`, +`kill(signal)`, `terminate(graceMs)`, `wait`, and `close`. It owns the process: +`close()`, a dropped connection, or exiting your program all terminate the +command rather than leaving it running in the box, and sessions cannot be +reattached. + +```ts +const dev = await box.exec.session({ cmd: "npm run dev", tty: true, rows: 24, cols: 80 }); +dev.write("rs\n"); // restart +dev.resize(50, 120); +dev.terminate(2000); // SIGTERM, then SIGKILL after 2s +``` + ### Files ```ts diff --git a/packages/sdk/package.json b/packages/sdk/package.json index 75a8193..9764e1a 100644 --- a/packages/sdk/package.json +++ b/packages/sdk/package.json @@ -45,8 +45,7 @@ "@types/ws": "^8.18.1", "prettier": "^3.8.1", "tsx": "^4.7.0", - "typescript": "^5.3.0", - "ws": "^8.20.0" + "typescript": "^5.3.0" }, "publishConfig": { "access": "public" @@ -59,6 +58,7 @@ }, "homepage": "https://upstash.com/docs/box", "dependencies": { + "ws": "^8.20.0", "zod-to-json-schema": "^3.25.1" }, "peerDependencies": { diff --git a/packages/sdk/src/__tests__/box-exec-session.test.ts b/packages/sdk/src/__tests__/box-exec-session.test.ts new file mode 100644 index 0000000..31079c9 --- /dev/null +++ b/packages/sdk/src/__tests__/box-exec-session.test.ts @@ -0,0 +1,375 @@ +import { describe, it, expect, afterEach, vi } from "vitest"; +import { WebSocketServer, type WebSocket as WsSocket } from "ws"; +import { Box } from "../client.js"; +import { mockResponse } from "./helpers.js"; + +const b64 = (s: string) => Buffer.from(s).toString("base64"); + +/** Start a mock exec-session WebSocket server on an ephemeral port. */ +async function startMockExecServer( + onConnection: (ws: WsSocket) => void, +): Promise<{ wss: WebSocketServer; port: number }> { + const wss = new WebSocketServer({ port: 0 }); + await new Promise((r) => wss.once("listening", () => r())); + wss.on("connection", onConnection); + const port = (wss.address() as { port: number }).port; + return { wss, port }; +} + +/** A Box whose baseUrl points at the local mock server. */ +async function boxForPort(port: number, timeout?: number): Promise { + const fetchMock = vi + .fn() + .mockResolvedValueOnce(mockResponse({ id: "box-123", status: "running", runtime: "node" })) + // Any later HTTP call (e.g. cd() -> POST /exec) succeeds; sessions use ws, not fetch. + .mockResolvedValue(mockResponse({ exit_code: 0 })); + vi.stubGlobal("fetch", fetchMock); + return Box.get("box-123", { apiKey: "k", baseUrl: `http://127.0.0.1:${port}`, timeout }); +} + +/** Server that replies `started` then `exit 0` to any start frame. */ +function trivialStart(onStart?: (frame: Record) => void) { + return (ws: WsSocket) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") { + onStart?.(f); + ws.send(JSON.stringify({ type: "started", pid: 1, execId: "e" })); + ws.send(JSON.stringify({ type: "exit", code: 0 })); + } + }); + }; +} + +describe("Box exec.session (WebSocket)", () => { + afterEach(() => { + vi.unstubAllGlobals(); + vi.restoreAllMocks(); + }); + + it("runs a non-TTY command: started -> stdout -> exit code", async () => { + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") { + ws.send(JSON.stringify({ type: "started", pid: 123, execId: "e1" })); + ws.send(JSON.stringify({ type: "stdout", data: b64("hello\n") })); + ws.send(JSON.stringify({ type: "exit", code: 7 })); + } + }); + }); + try { + const box = await boxForPort(port); + let out = ""; + const session = await box.exec.session({ + cmd: "echo hello", + onStdout: (d) => (out += Buffer.from(d).toString()), + }); + expect(session.pid).toBe(123); + expect(session.execId).toBe("e1"); + expect(await session.wait()).toBe(7); + expect(out).toBe("hello\n"); + } finally { + wss.close(); + } + }); + + it("round-trips stdin and terminates on a signal frame", async () => { + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") ws.send(JSON.stringify({ type: "started", pid: 5, execId: "e2" })); + else if (f.type === "stdin") + ws.send(JSON.stringify({ type: "stdout", data: f.data })); // echo + else if (f.type === "signal") ws.send(JSON.stringify({ type: "exit", code: 137 })); + }); + }); + try { + const box = await boxForPort(port); + let out = ""; + const s = await box.exec.session({ + argv: ["cat"], + onStdout: (d) => (out += Buffer.from(d).toString()), + }); + s.write("ping"); + await new Promise((r) => setTimeout(r, 50)); + expect(out).toBe("ping"); + s.kill("KILL"); + expect(await s.wait()).toBe(137); + } finally { + wss.close(); + } + }); + + it("sends argv (over cmd), tty, and dimensions in the start frame", async () => { + let startFrame: Record | undefined; + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") { + startFrame = f; + ws.send(JSON.stringify({ type: "started", pid: 1, execId: "e" })); + ws.send(JSON.stringify({ type: "exit", code: 0 })); + } + }); + }); + try { + const box = await boxForPort(port); + const s = await box.exec.session({ + cmd: "ignored", + argv: ["ls", "-la"], + tty: true, + rows: 24, + cols: 80, + }); + await s.wait(); + expect(startFrame?.argv).toEqual(["ls", "-la"]); + expect(startFrame?.cmd).toBeUndefined(); + expect(startFrame?.tty).toBe(true); + expect(startFrame?.rows).toBe(24); + expect(startFrame?.cols).toBe(80); + } finally { + wss.close(); + } + }); + + it("rejects when the server sends an error frame before start, and hangs up", async () => { + let closed!: () => void; + const serverSawClose = new Promise((r) => (closed = r)); + const { wss, port } = await startMockExecServer((ws) => { + ws.on("close", () => closed()); + ws.on("message", () => ws.send(JSON.stringify({ type: "error", message: "boom" }))); + }); + try { + const box = await boxForPort(port); + await expect(box.exec.session({ cmd: "x" })).rejects.toThrow(/boom/); + // A rejected session must not leak its socket. + await serverSawClose; + } finally { + wss.close(); + } + }); + + it("ends wait() and hangs up on an error frame after start", async () => { + let closed!: () => void; + const serverSawClose = new Promise((r) => (closed = r)); + const { wss, port } = await startMockExecServer((ws) => { + ws.on("close", () => closed()); + ws.on("message", (raw) => { + if (JSON.parse(raw.toString()).type === "start") { + ws.send(JSON.stringify({ type: "started", pid: 5, execId: "e" })); + ws.send(JSON.stringify({ type: "error", message: "late boom" })); + } + }); + }); + try { + const box = await boxForPort(port); + const session = await box.exec.session({ cmd: "x" }); + expect(await session.wait()).toBe(-1); + // Once wait() settles the caller considers the session over, so leaving + // the connection open would keep the process alive in the box. + await serverSawClose; + } finally { + wss.close(); + } + }); + + it.each([ + ["a zero pid", { type: "started", pid: 0, execId: "e" }], + ["no pid at all", { type: "started", execId: "e" }], + ])("rejects a started frame with %s", async (_label, startedFrame) => { + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + if (JSON.parse(raw.toString()).type === "start") ws.send(JSON.stringify(startedFrame)); + }); + }); + try { + const box = await boxForPort(port); + // A handle whose kill()/terminate() cannot reach the process is worse + // than no handle, so the handshake fails instead. + await expect(box.exec.session({ cmd: "x" })).rejects.toThrow(/without a usable pid/); + } finally { + wss.close(); + } + }); + + it.each(["stdout", "stderr"] as const)( + "contains a throwing on%s callback instead of crashing the process", + async (stream) => { + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + if (JSON.parse(raw.toString()).type === "start") { + ws.send(JSON.stringify({ type: "started", pid: 3, execId: "e" })); + ws.send(JSON.stringify({ type: stream, data: b64("boom\n") })); + } + }); + }); + try { + const box = await boxForPort(port); + const thrower = () => { + throw new Error("callback blew up"); + }; + const session = await box.exec.session({ + cmd: "x", + ...(stream === "stdout" ? { onStdout: thrower } : { onStderr: thrower }), + }); + // The throw must not escape the ws listener as an uncaught exception; + // the session ends instead so the host process survives. + expect(await session.wait()).toBe(-1); + } finally { + wss.close(); + } + }, + ); + + it("rejects locally (no socket) when neither cmd nor argv is given", async () => { + const box = await boxForPort(0); // port unused; must reject before connecting + await expect(box.exec.session({})).rejects.toThrow(/requires cmd or argv/); + await expect(box.exec.session({ cmd: "" })).rejects.toThrow(/requires cmd or argv/); + }); + + it("defaults cwd to the box cwd (honoring cd) and resolves an explicit cwd", async () => { + let frame: Record | undefined; + const { wss, port } = await startMockExecServer(trivialStart((f) => (frame = f))); + try { + const box = await boxForPort(port); + await box.cd("src"); + await (await box.exec.session({ cmd: "pwd" })).wait(); + expect(frame?.cwd).toBe("/workspace/home/src"); + + await (await box.exec.session({ cmd: "pwd", cwd: "nested" })).wait(); + expect(frame?.cwd).toBe("/workspace/home/src/nested"); + } finally { + wss.close(); + } + }); + + it("delivers stderr and forwards resize frames", async () => { + let resizeFrame: Record | undefined; + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") { + ws.send(JSON.stringify({ type: "started", pid: 1, execId: "e" })); + ws.send(JSON.stringify({ type: "stderr", data: b64("oops\n") })); + } else if (f.type === "resize") { + resizeFrame = f; + ws.send(JSON.stringify({ type: "exit", code: 0 })); + } + }); + }); + try { + const box = await boxForPort(port); + let err = ""; + const s = await box.exec.session({ + tty: true, + cmd: "x", + onStderr: (d) => (err += Buffer.from(d).toString()), + }); + await new Promise((r) => setTimeout(r, 30)); + expect(err).toBe("oops\n"); + s.resize(30, 100); + await s.wait(); + expect(resizeFrame).toMatchObject({ rows: 30, cols: 100 }); + } finally { + wss.close(); + } + }); + + it("close() ends the session (wait resolves)", async () => { + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") ws.send(JSON.stringify({ type: "started", pid: 1, execId: "e" })); + // never sends exit; client close() must settle wait() + }); + }); + try { + const box = await boxForPort(port); + const s = await box.exec.session({ cmd: "sleep 999" }); + s.close(); + expect(await s.wait()).toBe(-1); + } finally { + wss.close(); + } + }); + + it("kill() rejects a signal outside the allowlist", async () => { + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") ws.send(JSON.stringify({ type: "started", pid: 1, execId: "e" })); + }); + }); + try { + const box = await boxForPort(port); + const s = await box.exec.session({ cmd: "sleep 999" }); + expect(() => s.kill("BOGUS")).toThrow(/unsupported signal/); + expect(() => s.kill("SIGKILL")).not.toThrow(); + s.close(); + } finally { + wss.close(); + } + }); + + it("endStdin() sends a stdin_close frame; the process exits on EOF (no kill)", async () => { + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") ws.send(JSON.stringify({ type: "started", pid: 1, execId: "e" })); + else if (f.type === "stdin") ws.send(JSON.stringify({ type: "stdout", data: f.data })); + else if (f.type === "stdin_close") ws.send(JSON.stringify({ type: "exit", code: 0 })); + }); + }); + try { + const box = await boxForPort(port); + let out = ""; + const s = await box.exec.session({ + argv: ["cat"], + onStdout: (d) => (out += Buffer.from(d).toString()), + }); + s.write("hi\n"); + await new Promise((r) => setTimeout(r, 30)); + s.endStdin(); + expect(await s.wait()).toBe(0); + expect(out).toBe("hi\n"); + } finally { + wss.close(); + } + }); + + it("terminate() sends a terminate frame with graceMs", async () => { + let termFrame: Record | undefined; + const { wss, port } = await startMockExecServer((ws) => { + ws.on("message", (raw) => { + const f = JSON.parse(raw.toString()); + if (f.type === "start") ws.send(JSON.stringify({ type: "started", pid: 1, execId: "e" })); + else if (f.type === "terminate") { + termFrame = f; + ws.send(JSON.stringify({ type: "exit", code: 143 })); + } + }); + }); + try { + const box = await boxForPort(port); + const s = await box.exec.session({ cmd: "sleep 999" }); + s.terminate(2000); + expect(await s.wait()).toBe(143); + expect(termFrame).toMatchObject({ type: "terminate", graceMs: 2000 }); + } finally { + wss.close(); + } + }); + + it("times out the handshake when the server never sends started", async () => { + const { wss, port } = await startMockExecServer(() => { + /* accept the socket but never reply */ + }); + try { + const box = await boxForPort(port, 150); + await expect(box.exec.session({ cmd: "x" })).rejects.toThrow(/handshake timed out/); + } finally { + wss.close(); + } + }); +}); diff --git a/packages/sdk/src/__tests__/integration/exec-session.integration.test.ts b/packages/sdk/src/__tests__/integration/exec-session.integration.test.ts new file mode 100644 index 0000000..cd8cea0 --- /dev/null +++ b/packages/sdk/src/__tests__/integration/exec-session.integration.test.ts @@ -0,0 +1,226 @@ +import { describe, it, expect, beforeAll, afterAll } from "vitest"; +import { Agent, Box, ClaudeCode } from "../../index.js"; +import { UPSTASH_BOX_API_KEY } from "./setup.js"; + +const text = (b: Uint8Array) => Buffer.from(b).toString(); + +describe.skipIf(!UPSTASH_BOX_API_KEY)("exec.session", () => { + let box: Box; + + beforeAll(async () => { + box = await Box.create({ + apiKey: UPSTASH_BOX_API_KEY!, + agent: { runner: Agent.ClaudeCode, model: ClaudeCode.Opus_4_6 }, + }); + }, 120000); + + afterAll(async () => { + try { + await box?.delete(); + } catch { + // cleanup best-effort + } + }, 30000); + + it("streams stdout and stderr separately and reports the exit code", async () => { + let out = ""; + let err = ""; + const session = await box.exec.session({ + argv: ["sh", "-c", "echo to-stdout; echo to-stderr 1>&2; exit 42"], + onStdout: (b) => (out += text(b)), + onStderr: (b) => (err += text(b)), + }); + + expect(session.pid).toBeGreaterThan(0); + expect(session.execId).not.toBe(""); + expect(await session.wait()).toBe(42); + expect(out.trim()).toBe("to-stdout"); + expect(err.trim()).toBe("to-stderr"); + }); + + it("runs argv without a shell", async () => { + let out = ""; + const session = await box.exec.session({ + argv: ["echo", "$HOME; rm -rf /"], + onStdout: (b) => (out += text(b)), + }); + await session.wait(); + // A shell would expand $HOME and treat `;` as a separator. + expect(out.trim()).toBe("$HOME; rm -rf /"); + }); + + it("runs cmd through a shell", async () => { + let out = ""; + const session = await box.exec.session({ + cmd: "echo shell-$((1+1))", + onStdout: (b) => (out += text(b)), + }); + await session.wait(); + expect(out.trim()).toBe("shell-2"); + }); + + it("writes stdin and lets endStdin() finish an EOF-reading command", async () => { + let out = ""; + const session = await box.exec.session({ + argv: ["sort"], + onStdout: (b) => (out += text(b)), + }); + session.write("banana\napple\ncherry\n"); + session.endStdin(); + + expect(await session.wait()).toBe(0); + expect(out).toBe("apple\nbanana\ncherry\n"); + }); + + it("honors cwd and overlays env", async () => { + await box.files.mkdir("session-proj/src", { parents: true }); + let out = ""; + const session = await box.exec.session({ + argv: ["sh", "-c", "pwd; echo $MY_VAR"], + cwd: "session-proj/src", + env: ["MY_VAR=from-test"], + onStdout: (b) => (out += text(b)), + }); + await session.wait(); + + expect(out).toContain("/workspace/home/session-proj/src"); + expect(out).toContain("from-test"); + await box.files.remove("session-proj", { recursive: true }); + }); + + it("drops blocked env keys but passes ordinary ones", async () => { + let out = ""; + const session = await box.exec.session({ + argv: ["sh", "-c", "echo LD=[$LD_PRELOAD] SAFE=[$SAFE]"], + env: ["LD_PRELOAD=/tmp/evil.so", "SAFE=yes"], + onStdout: (b) => (out += text(b)), + }); + await session.wait(); + + expect(out).toContain("LD=[]"); + expect(out).toContain("SAFE=[yes]"); + }); + + it("terminate() stops a long-running process", async () => { + const session = await box.exec.session({ argv: ["sleep", "300"] }); + session.terminate(1000); + expect(await session.wait()).not.toBe(0); + }); + + it("kill() reaps the whole process tree", async () => { + const running = async () => + ( + await box.exec.command( + `c=0; for d in /proc/[0-9]*; do [ "$(cat "$d/comm" 2>/dev/null)" = "sleep" ] && grep -qs 4711 "$d/cmdline" && c=$((c+1)); done; echo $c`, + ) + ).result.trim(); + + const session = await box.exec.session({ cmd: "sleep 4711 & sleep 4712 & wait" }); + await new Promise((r) => setTimeout(r, 800)); + expect(await running()).not.toBe("0"); + + session.kill("TERM"); + await session.wait(); + await new Promise((r) => setTimeout(r, 500)); + expect(await running()).toBe("0"); + }); + + it("allocates a real PTY at the requested size and accepts interactive input", async () => { + let out = ""; + const session = await box.exec.session({ + tty: true, + rows: 24, + cols: 80, + cmd: "tty; stty size; read line; echo GOT=$line; exit 0", + onStdout: (b) => (out += text(b)), + }); + await new Promise((r) => setTimeout(r, 400)); + session.write("hello-pty\n"); + + expect(await session.wait()).toBe(0); + expect(out).toContain("/dev/pts/"); + // Size must be right from the first read, not applied after the process starts. + expect(out).toContain("24 80"); + expect(out).toContain("GOT=hello-pty"); + }); + + it("keeps a long-lived process for multiple round-trips", async () => { + let out = ""; + const session = await box.exec.session({ + argv: ["cat"], + onStdout: (b) => (out += text(b)), + }); + for (const msg of ["req-1\n", "req-2\n", "req-3\n"]) { + session.write(msg); + await new Promise((r) => setTimeout(r, 150)); + } + expect(out).toContain("req-1"); + expect(out).toContain("req-2"); + expect(out).toContain("req-3"); + + session.kill("KILL"); + await session.wait(); + }); + + it("runs sessions concurrently without crosstalk", async () => { + const results = await Promise.all( + [1, 2, 3, 4].map(async (n) => { + let out = ""; + const session = await box.exec.session({ + argv: ["sh", "-c", `sleep 0.${n}; echo worker-${n}`], + onStdout: (b) => (out += text(b)), + }); + return { n, code: await session.wait(), out: out.trim() }; + }), + ); + for (const r of results) { + expect(r.code).toBe(0); + expect(r.out).toBe(`worker-${r.n}`); + } + }); + + it("does not leak env between sessions", async () => { + await (await box.exec.session({ argv: ["sh", "-c", "export LEAK=nope; true"] })).wait(); + let out = ""; + const session = await box.exec.session({ + argv: ["sh", "-c", "echo LEAK=[$LEAK]"], + onStdout: (b) => (out += text(b)), + }); + await session.wait(); + expect(out).toContain("LEAK=[]"); + }); + + it("close() ends the session and stops the process", async () => { + const session = await box.exec.session({ argv: ["sleep", "600"] }); + const pid = session.pid; + session.close(); + await session.wait(); + + await new Promise((r) => setTimeout(r, 1500)); + const alive = ( + await box.exec.command(`[ -d /proc/${pid} ] && echo yes || echo no`) + ).result.trim(); + expect(alive).toBe("no"); + }); + + it("rejects an empty command locally and an unsupported signal", async () => { + await expect(box.exec.session({})).rejects.toThrow(/requires cmd or argv/); + + const session = await box.exec.session({ argv: ["sleep", "60"] }); + expect(() => session.kill("BOGUS")).toThrow(/unsupported signal/); + session.close(); + await session.wait(); + }); + + it("makes session writes visible to the files API", async () => { + const session = await box.exec.session({ + argv: ["sh", "-c", "echo written-by-session > session-out.txt"], + }); + expect(await session.wait()).toBe(0); + + expect((await box.files.read("session-out.txt")).trim()).toBe("written-by-session"); + const stat = await box.files.stat("session-out.txt"); + expect(stat.type).toBe("file"); + await box.files.remove("session-out.txt"); + }); +}); diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index 78bf580..d6e6768 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -22,6 +22,8 @@ import { type ErrorResponse, type FileEntry, type FileStat, + type ExecSessionOptions, + type ExecSessionHandle, type GitCloneOptions, type GitExecOptions, type GitExecResult, @@ -153,6 +155,31 @@ function safeExecOutputLength(buffer: string): number { return buffer.length; } +/** + * Minimal structural type for a `ws` WebSocket, so the public API and its + * generated `.d.ts` do not depend on `@types/ws`. + */ +interface WsLike { + send(data: string): void; + close(): void; + on(event: "open", cb: () => void): void; + on(event: "message", cb: (data: { toString(): string }) => void): void; + on(event: "close", cb: () => void): void; + on(event: "error", cb: (err: Error) => void): void; +} + +/** Signals an exec session accepts, mirroring the backend allowlist. */ +const EXEC_SESSION_SIGNALS = new Set([ + "TERM", + "KILL", + "INT", + "HUP", + "TSTP", + "QUIT", + "USR1", + "USR2", +]); + /** * Error thrown by the Box SDK */ @@ -638,6 +665,12 @@ export class Box { code: (options: CodeExecutionOptions) => Promise>; stream: (command: string) => Promise>; streamCode: (options: CodeExecutionOptions) => Promise>; + /** + * Open a live, interactive command session over a WebSocket: stdin, + * streamed stdout/stderr, resize, and signals. Node-only. See + * {@link ExecSessionOptions}. + */ + session: (options: ExecSessionOptions) => Promise; }; /** Schedule operations namespace */ @@ -821,6 +854,7 @@ export class Box { code: (options) => this._execCode(options), stream: (command) => this._execStream(command), streamCode: (options) => this._execStreamCode(options), + session: (options) => this._execSession(options), }; this.files = { @@ -1793,6 +1827,180 @@ export class Box { return run; } + /** + * Open a live, interactive command session over a WebSocket. Resolves once the + * process has started (so {@link ExecSessionHandle.pid} is available); output + * flows to the `onStdout`/`onStderr` callbacks. Node-only: the auth header is + * set on the WebSocket handshake, which browsers cannot do. + */ + private async _execSession(options: ExecSessionOptions): Promise { + const hasArgv = !!options.argv && options.argv.length > 0; + if (!hasArgv && !options.cmd) { + throw new BoxError("exec.session requires cmd or argv"); + } + + const start: Record = { type: "start" }; + if (hasArgv) start.argv = options.argv; + else start.cmd = options.cmd; + if (options.tty) start.tty = true; + // Default to the box's current directory (honoring cd()), matching exec.command. + start.cwd = options.cwd ? this._resolvePath(options.cwd) : this._cwd; + if (options.rows) start.rows = options.rows; + if (options.cols) start.cols = options.cols; + if (options.env) start.env = options.env; + + const wsModule = await import("ws"); + const WsCtor = (wsModule.default ?? + (wsModule as unknown as { WebSocket: unknown }).WebSocket) as new ( + url: string, + opts: { headers: Record }, + ) => WsLike; + + const wsUrl = this._baseUrl.replace(/^http/, "ws") + `/v2/box/${this.id}/exec-session`; + const socket = new WsCtor(wsUrl, { headers: { ...this._headers } }); + + let resolveExit!: (code: number) => void; + const exitPromise = new Promise((r) => (resolveExit = r)); + let exited = false; + const settleExit = (code: number) => { + if (!exited) { + exited = true; + resolveExit(code); + } + }; + + const send = (obj: unknown) => socket.send(JSON.stringify(obj)); + + // Output callbacks run inside the socket's "message" listener, so a throw + // would escape the emitter as an uncaught exception and could take down the + // host process. End the session instead, matching the Python handles. + const dispatch = (cb: ((data: Uint8Array) => void) | undefined, data: unknown) => { + if (!cb) return; + try { + cb(new Uint8Array(Buffer.from(String(data), "base64"))); + } catch { + settleExit(-1); + socket.close(); + } + }; + + return await new Promise((resolve, reject) => { + let started = false; + const timer = setTimeout(() => { + if (!started) { + reject(new BoxError("exec.session handshake timed out")); + try { + socket.close(); + } catch { + // already closing + } + } + }, this._timeout); + const failStart = (message: string) => { + if (!started) { + clearTimeout(timer); + reject(new BoxError(message)); + } + }; + + socket.on("open", () => send(start)); + socket.on("error", (err: Error) => { + failStart(`exec-session connection failed: ${err.message}`); + settleExit(-1); + }); + socket.on("close", () => { + failStart("exec-session closed before start"); + settleExit(-1); + }); + socket.on("message", (raw) => { + let frame: Record; + try { + frame = JSON.parse(raw.toString()) as Record; + } catch { + return; + } + switch (frame.type) { + case "started": { + // A session whose process cannot be signaled is useless, so a + // missing or non-positive pid fails the handshake rather than + // handing back a handle whose kill()/terminate() go nowhere. + const pid = typeof frame.pid === "number" ? frame.pid : 0; + if (pid <= 0) { + failStart("exec-session started without a usable pid"); + try { + socket.close(); + } catch { + // already closing + } + break; + } + started = true; + clearTimeout(timer); + const handle: ExecSessionHandle = { + pid, + execId: typeof frame.execId === "string" ? frame.execId : "", + write: (data) => { + if (exited) return; + const bytes = + typeof data === "string" ? Buffer.from(data, "utf8") : Buffer.from(data); + send({ type: "stdin", data: bytes.toString("base64") }); + }, + endStdin: () => { + if (exited) return; + send({ type: "stdin_close" }); + }, + resize: (rows, cols) => { + if (exited) return; + send({ type: "resize", rows, cols }); + }, + kill: (signal) => { + if (exited) return; + const sig = (signal ?? "TERM").trim().toUpperCase().replace(/^SIG/, ""); + if (!EXEC_SESSION_SIGNALS.has(sig)) { + throw new BoxError(`unsupported signal: ${signal}`); + } + send({ type: "signal", signal: sig }); + }, + terminate: (graceMs) => { + if (exited) return; + send({ type: "terminate", ...(graceMs && graceMs > 0 ? { graceMs } : {}) }); + }, + wait: () => exitPromise, + close: () => { + try { + socket.close(); + } catch { + // already closing + } + }, + }; + resolve(handle); + break; + } + case "stdout": + dispatch(options.onStdout, frame.data); + break; + case "stderr": + dispatch(options.onStderr, frame.data); + break; + case "exit": + settleExit(typeof frame.code === "number" ? frame.code : -1); + socket.close(); + break; + case "error": + // Before start: reject session(). After start: end wait() with -1. + if (!started) failStart(`exec-session error: ${String(frame.message)}`); + else settleExit(-1); + // Hang up either way. A rejected session must not leak its socket, + // and once wait() has settled the caller considers the session over, + // so leaving the connection open would keep the process alive. + socket.close(); + break; + } + }); + }); + } + /** * Stream output from inline code execution in the box. */ @@ -3028,6 +3236,12 @@ export class EphemeralBox { code: (options: CodeExecutionOptions) => Promise>; stream: (command: string) => Promise>; streamCode: (options: CodeExecutionOptions) => Promise>; + /** + * Open a live, interactive command session over a WebSocket: stdin, + * streamed stdout/stderr, resize, and signals. Node-only. See + * {@link ExecSessionOptions}. + */ + session: (options: ExecSessionOptions) => Promise; }; /** Schedule operations namespace */ diff --git a/packages/sdk/src/index.ts b/packages/sdk/src/index.ts index 6fb2d2b..48369f6 100644 --- a/packages/sdk/src/index.ts +++ b/packages/sdk/src/index.ts @@ -50,6 +50,8 @@ export type { UploadFileEntry, FileEntry, FileStat, + ExecSessionOptions, + ExecSessionHandle, GitCloneOptions, GitExecOptions, GitExecResult, diff --git a/packages/sdk/src/types.ts b/packages/sdk/src/types.ts index fa79060..78af59d 100644 --- a/packages/sdk/src/types.ts +++ b/packages/sdk/src/types.ts @@ -899,6 +899,79 @@ export interface FileEntry { mod_time: string; } +/** + * Options for {@link Box} `exec.session()` — one live, interactive command over + * a WebSocket. Provide `argv` for an exact program (no shell) or `cmd` for a + * `bash -lc` string. This is a Node-only API (it sets an auth header on the + * WebSocket handshake, which browsers cannot do). + */ +export interface ExecSessionOptions { + /** Command run via `bash -lc`. Ignored when `argv` is set. */ + cmd?: string; + /** Exact program and arguments, run without a shell. Takes precedence over `cmd`. */ + argv?: string[]; + /** Allocate a PTY: resizable, with stdout and stderr merged onto stdout. */ + tty?: boolean; + /** Working directory; defaults to the box workspace. */ + cwd?: string; + /** Initial PTY row count (TTY sessions). */ + rows?: number; + /** Initial PTY column count (TTY sessions). */ + cols?: number; + /** Extra environment entries as `KEY=VALUE`, overlaid on the box environment. */ + env?: string[]; + /** Called with decoded stdout bytes as they arrive. */ + onStdout?: (data: Uint8Array) => void; + /** Called with decoded stderr bytes as they arrive (non-TTY sessions). */ + onStderr?: (data: Uint8Array) => void; +} + +/** + * A live command session. `session()` resolves this once the process has + * started; output flows to the `onStdout`/`onStderr` callbacks passed to + * `session()`. + * + * The session owns the process: losing the connection kills it. `close()`, a + * dropped network link, or exiting the program all terminate the command rather + * than leaving it running in the box, and sessions cannot be reattached. Use + * `wait()` to run something to completion. + */ +export interface ExecSessionHandle { + /** + * In-box (container-namespace) PID of the running process. Always non-zero, + * so there is no need to guard on this: the server fails the handshake rather + * than starting a session it cannot signal, and `session()` rejects a + * `started` frame that carries no usable pid. + */ + readonly pid: number; + /** Server-side exec id. */ + readonly execId: string; + /** Write bytes to the process stdin. */ + write(data: string | Uint8Array): void; + /** + * Close stdin (send EOF). A command that reads until EOF (e.g. `cat`, `sort`) + * then exits on its own; stdout/stderr keep flowing until it does. + */ + endStdin(): void; + /** Resize the PTY (TTY sessions). */ + resize(rows: number, cols: number): void; + /** Send a signal to the process tree. Defaults to `TERM`. */ + kill(signal?: string): void; + /** + * Graceful stop driven server-side: SIGTERM now, then SIGKILL after `graceMs` + * (default is the server's grace) if the process has not exited. + * + * Only the first call starts the sequence; later ones are ignored, so the + * grace cannot be extended or shortened once it is running. Send + * `kill("KILL")` to stop the process immediately instead. + */ + terminate(graceMs?: number): void; + /** Resolve with the process exit code (`-1` if still running after a forced teardown). */ + wait(): Promise; + /** Close the connection, terminating the process if still running. */ + close(): void; +} + /** * Filesystem metadata for a single path, returned by `files.stat`. * diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 252e7e1..5150d82 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -67,6 +67,9 @@ importers: packages/sdk: dependencies: + ws: + specifier: ^8.20.0 + version: 8.20.0 zod: specifier: ^3.25.0 || ^4.0.0 version: 4.3.6 @@ -89,9 +92,6 @@ importers: typescript: specifier: ^5.3.0 version: 5.9.3 - ws: - specifier: ^8.20.0 - version: 8.20.0 packages: