diff --git a/CHANGELOG.md b/CHANGELOG.md index 8633acd..7782e41 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,12 @@ below corresponds to a tag of the same name. ### Added +- **exec:** `colab exec --json` prints one JSON object per executed cell on + stdout (`status`, `stdout`, `stderr`, `data`, `images`, `error`) instead of + streaming output, for scripts and AI agents. Tracebacks are plain text + without ANSI colour codes, images are saved to disk and reported by path, + `[colab]` progress lines go to stderr, and the exit status is 1 when any + cell failed. Without `--json`, output and exit status are unchanged. - **ssh:** `colab ssh [-s NAME] -- CMD` runs one command on an existing runtime and exits with its status. Output is byte-exact when piped (a PTY only when stdin and stdout are terminals), and it never creates a runtime. Its handshake is tagged `ssh-mode-cmd`. ### Removed diff --git a/README.md b/README.md index 53c16ab..7dd11a0 100644 --- a/README.md +++ b/README.md @@ -79,7 +79,7 @@ Run `colab --help` to view specific options, defaults, and detailed he | Command | Description | | --- | --- | | `colab run [--gpu GPU] [--tpu TPU] [--high-mem] [--keep] SCRIPT [ARGS...]` | Run a local script on a fresh VM, forwarding arguments, then release it | -| `colab exec [-s NAME] [-f FILE] [--output-image PATH]` | Execute Python code from stdin, a local `.py` file, or a `.ipynb` notebook | +| `colab exec [-s NAME] [-f FILE] [--output-image PATH] [--json]` | Execute Python code from stdin, a local `.py` file, or a `.ipynb` notebook (`--json` prints one machine-readable result per cell) | | `colab repl [-s NAME] [--output-image PATH]` | Start an interactive Python REPL on the VM (exits cleanly on piped EOF) | | `colab ssh [-s NAME] [--proxy-mode] [-i KEY] [--gpu GPU] [--tpu TPU] [--high-mem] [-- CMD ...]` | Open an SSH shell to the runtime over WebSocket, run a one-shot remote command (`-- CMD`), or act as an OpenSSH `ProxyCommand` bridge for IDE remote-dev | diff --git a/docs/02_execution_and_interactive.md b/docs/02_execution_and_interactive.md index f8bbb18..14c146b 100644 --- a/docs/02_execution_and_interactive.md +++ b/docs/02_execution_and_interactive.md @@ -9,6 +9,8 @@ log: 2026-10-24: `--timeout` for `colab exec` (and `colab run`) is now an idle limit: execution fails only after that many seconds without output, so code that keeps printing is never cut off, and `--timeout 0` means no limit. On 0.8.0/0.9.0 the old deadline was never enforced, so code that is silent for longer than `--timeout` now fails where it used to complete. `colab repl` and `colab install`, which pass no timeout, wait until execution finishes, or until the kernel dies or restarts. None of them busy-loop while waiting. `colab auth` and `colab drivemount` are unchanged. See #183. 2026-10-29: Removed the `colab console` subcommand and its `/colab/tty` transport (`console.py`, its tests, and `integration/repro_piped_console/`). Use `colab ssh -s NAME` for an interactive shell and `colab ssh -s NAME -- CMD` for one-shot commands, or pipe `!`-prefixed shell commands to `colab exec` where ssh can't work. + +2026-11-11: Added `colab exec --json`. Each executed cell becomes one JSON object on stdout (NDJSON) instead of streamed text, `[colab]` progress lines move to stderr, and the exit status is 1 when any cell failed. Without the flag, `colab exec` is unchanged (it still exits 0 after a kernel error). --- # Design: Execution and Interactive Interaction (`repl`, `exec`) @@ -33,6 +35,11 @@ Execution involves sending Python code (or shell commands) to the Jupyter kernel - If file path is remote: Execute `!python `. - **Multi-Modal Output**: Handle `display_data` messages (e.g., `image/png`, `text/html`). For the CLI, we'll save images to temporary files and print their paths, or if the terminal supports it (e.g., iTerm2), inline them. - **Timeout Configuration**: Exposes a `--timeout` flag (default 30s): the longest silence allowed before execution fails. Raise it so long-running silent tasks (like model compilation or data downloading) are not prematurely killed; `0` means no limit. +- **Machine-Readable Output (`--json`)**: For scripts and AI agents that must decide what to do next from the result. Plain `colab exec` exits 0 after a kernel error and prints an ANSI-coloured traceback, so a caller can only guess at failure by scraping text. With `--json`: + - stdout carries one JSON object per executed cell, one per line (NDJSON), written when the cell finishes. Nothing else is written to stdout; `[colab]` progress lines go to stderr. + - Fields: `cell_index`, `cell_id` (notebook cell id, else `null`), `status` (`"ok"` or `"error"`), `stdout`, `stderr` (the kernel's streams, concatenated), `data` (one object per `execute_result` / `display_data` output: its `output_type` plus its non-image MIME entries such as `text/plain`), `images` (paths of the PNG/JPEG outputs saved to disk, honouring `--output-image`; never inline base64), and `error` (`null`, or `ename`, `evalue` and `traceback` as plain text with ANSI escapes removed). + - The exit status is 1 when any cell has `status: "error"`. Notebook cells after a failing cell still run, as without the flag. + - A client-side failure (idle `--timeout`, lost connection) is reported as a final record whose `error.ename` is the exception type, before the command fails as it does without the flag. ## Implementation Details - **Kernel Management**: `ColabRuntime` (from `colab-agent`) already handles message signing and message types. @@ -48,6 +55,7 @@ TDD is mandatory for all execution features. - **Test Case**: Verify `display_data` (specifically `image/png`) triggers the correct local handling (saving or display). ### 2. TTY and Piping +- **Test Case**: `tests/test_exec_json.py` covers the `--json` record for success, kernel error (ANSI stripped, exit 1), images (saved, path reported, no base64 on stdout), per-cell notebook records, client-side failures, a lost session, and that the default mode still exits 0 on a kernel error. - **Test Case**: Mock `sys.stdin.isatty()` to verify `colab repl` correctly switches between interactive mode and one-shot piped execution. - **Test Case**: Verify large piped inputs are handled without buffer overflow or truncation. - **Test Case**: `print_kitty` is a no-op when `sys.stdout.isatty()` is false (regression: previously emitted ANSI/base64 into pipes and files). diff --git a/skills/colab-operator/SKILL.md b/skills/colab-operator/SKILL.md index fe41419..1bc5b7a 100644 --- a/skills/colab-operator/SKILL.md +++ b/skills/colab-operator/SKILL.md @@ -52,6 +52,7 @@ by running `uv tool install google-colab-cli` or `pip install google-colab-cli`. - **Preferred**: `colab exec -s -f ` runs a local script on the remote VM (read locally, sent to the kernel — no manual upload needed). - **Piped code**: `echo "print(1)" | colab exec -s ` or `cat script.py | colab exec -s `. - **Notebooks**: `colab exec -s -f nb.ipynb` runs each code cell and writes results to `_output.ipynb` next to the input. A `# @title Foo` first line labels the cell in progress output. +- **Check the result with `--json`**: plain `colab exec` exits 0 even when the code raises, so don't infer success from the exit status. `colab exec -s --json` prints one JSON object per executed cell on stdout (`status` is `"ok"` or `"error"`; `stdout`, `stderr`; `data` with `text/plain` results; `images` as saved file paths; `error` with `ename`, `evalue` and a colour-free `traceback`) and exits 1 if any cell failed. `[colab]` progress lines go to stderr. Output arrives when a cell finishes, not live, so keep `--timeout` in mind for long cells. - **Plots/images**: PNG/JPEG outputs are intercepted. Use `--output-image ` on `exec`/`repl` to save to a known location (otherwise a temp path is printed). Inline terminal-image escapes are auto-suppressed when stdout isn't a TTY, so piped/captured output stays clean. - **Shell**: run a one-shot shell command with `colab ssh -s -- `: it runs in `/content` and exits with the command's status. Output is byte-exact when stdout is piped or redirected; ssh allocates a PTY only when stdin and stdout are both terminals, so a shell tool that runs commands in a terminal should add ` -- 'nvidia-smi | head -5'`); several arguments are an exact argv, so pass untrusted values as separate arguments after `--` where the program accepts it (e.g. `colab ssh -s -- ls -- "$f"`), never inside a shell line or a piped script. Pipe a multi-line script to `colab ssh -s -- bash`. Use `colab exec` for Python code. When ssh can't work (no local `ssh` client, no `~/.ssh/id_ed25519` or `id_ecdsa` key and no `-i KEY`, a tool not on the minimal ssh environment's `PATH` such as `gcloud` or `node`, or the runtime's limited concurrent SSH connections in use, which the server answers with 429, e.g. while a person or an IDE is connected), fall back to `exec` with IPython's `!` prefix: `echo '!df -h /content' | colab exec -s ` (single quotes stop your local shell expanding `!`). Each `!` line runs in its own shell with no input (a command that prompts hangs), and IPython evaluates `{expr}` as Python even inside quotes and expands `$name` outside single quotes. Start a multi-line script with a `%%bash` line (not with `--env`, which prepends Python lines), use non-interactive flags such as `apt-get -y`, and never put untrusted text in a `%%bash` script or a `!` line (even via `{var}` or `$var`): pass it with `--env NAME=VALUE` and run `subprocess.run(["cmd", "--", os.environ["NAME"]])` from Python. For an interactive shell on the VM, a human runs `colab ssh -s `. - **Never run `colab repl`, `colab ssh` without `-- CMD`, `colab auth`, or `colab drivemount` interactively from an agent** — `repl`, `auth` and `drivemount` expect a TTY and will hang or fail, and bare `colab ssh` with no saved session creates a billable runtime; an interactive shell is for a human. `repl` accepts piped stdin and exits on EOF; `auth`/`drivemount` genuinely require a human at the terminal. diff --git a/src/colab_cli/commands/execution.py b/src/colab_cli/commands/execution.py index a21d206..843353f 100644 --- a/src/colab_cli/commands/execution.py +++ b/src/colab_cli/commands/execution.py @@ -13,6 +13,7 @@ # limitations under the License. import datetime +import json import nbformat import os import re @@ -25,12 +26,20 @@ from typing_extensions import Annotated from colab_cli.runtime import ColabRuntime -from colab_cli.utils import handle_image, is_terminal_error, render_display_data +from colab_cli.utils import ( + handle_image, + is_terminal_error, + render_display_data, + save_image, +) _console = Console() TITLE_REGEX = re.compile(r"^\s*#\s*@title\s+(.*)", re.MULTILINE) ENV_KEY_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +# CSI escape sequences (colours, cursor moves) that IPython puts in tracebacks. +ANSI_ESCAPE_REGEX = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]") +IMAGE_MIME_TYPES = ("image/png", "image/jpeg") def is_stdin_tty() -> bool: @@ -142,6 +151,64 @@ def display_output(out, output_image=None): pass +def _error_record(ename: str, evalue: str, traceback: List[str]) -> dict: + return { + "ename": ename, + "evalue": evalue, + "traceback": ANSI_ESCAPE_REGEX.sub("", "\n".join(traceback)), + } + + +def build_json_record( + outputs: List[dict], + cell_index: int, + cell_id: Optional[str], + output_image: Optional[str] = None, +) -> dict: + """Summarize one cell's kernel outputs as a JSON-serializable record. + + This is the `colab exec --json` contract: streams are concatenated, images + are written to disk and reported by path (never inlined as base64), and + the traceback is plain text with ANSI colour codes removed. + """ + record = { + "cell_index": cell_index, + "cell_id": cell_id, + "status": "ok", + "stdout": "", + "stderr": "", + "data": [], + "images": [], + "error": None, + } + for out in outputs: + output_type = out.get("output_type") + if output_type == "stream": + name = "stderr" if out.get("name") == "stderr" else "stdout" + record[name] += out.get("text", "") + elif output_type == "error": + record["status"] = "error" + record["error"] = _error_record( + out.get("ename", "Error"), + out.get("evalue", ""), + out.get("traceback", []), + ) + elif "data" in out: + data = out["data"] + for mime_type in IMAGE_MIME_TYPES: + if mime_type in data: + record["images"].append( + save_image(data[mime_type], mime_type, output_image) + ) + break + rest = {k: v for k, v in data.items() if not k.startswith("image/")} + if rest: + record["data"].append( + {"output_type": output_type or "display_data", **rest} + ) + return record + + def exec_command( session: Annotated[ Optional[str], typer.Option("-s", "--session", help="Session name") @@ -169,10 +236,27 @@ def exec_command( ), ), ] = None, + json_output: Annotated[ + bool, + typer.Option( + "--json", + help=( + "Print one JSON object per executed cell on stdout instead of " + "streaming output, and exit 1 if any cell failed." + ), + ), + ] = False, ): """Execute code in a session""" from colab_cli.common import state + def info(message: str) -> None: + # With --json, stdout carries only JSON records. + if json_output: + typer.echo(message, err=True) + else: + typer.echo(message) + env_vars = _parse_env_vars(env) name = state.resolve_session(session) s = state.get_session(name) @@ -180,7 +264,7 @@ def exec_command( code_blocks = [] if file: if file.endswith(".ipynb"): - typer.echo(f"[colab] Parsing notebook '{file}'...") + info(f"[colab] Parsing notebook '{file}'...") with open(file, "r", encoding="utf-8") as f: nb = nbformat.read(f, as_version=4) for cell in nb.cells: @@ -227,13 +311,12 @@ def on_sess_started(sid): ) except Exception as e: if is_terminal_error(e): - typer.echo( - f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up." - ) + info(f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up.") state.prune_session(name) raise typer.Exit(1) raise e + failed = False try: is_nb = file and file.endswith(".ipynb") s.running = f"exec({file or 'stdin'})" @@ -252,7 +335,7 @@ def on_sess_started(sid): identifier = "" identifier_str = f" - {identifier}" if identifier else "" - typer.echo( + info( f"[colab] Executing cell {i + 1}/{len(code_blocks)}{identifier_str}..." ) @@ -263,11 +346,26 @@ def on_sess_started(sid): ) state.store.add(s) - outputs = runtime.execute_code( - code, - output_hook=lambda o: display_output(o, output_image), - timeout=timeout, - ) + if json_output: + try: + outputs = runtime.execute_code(code, timeout=timeout) + except Exception as e: + # Timeouts and lost connections never reach the kernel's + # outputs; report them in the same shape before failing. + record = build_json_record([], i, block.get("id")) + record["status"] = "error" + record["error"] = _error_record(type(e).__name__, str(e), []) + typer.echo(json.dumps(record)) + raise + record = build_json_record(outputs, i, block.get("id"), output_image) + failed = failed or record["status"] == "error" + typer.echo(json.dumps(record)) + else: + outputs = runtime.execute_code( + code, + output_hook=lambda o: display_output(o, output_image), + timeout=timeout, + ) if "cell" in block: save_output(outputs, block["cell"]) state.history.log_event( @@ -286,9 +384,11 @@ def on_sess_started(sid): runtime.stop() if file and file.endswith(".ipynb"): output_file = os.path.splitext(file)[0] + "_output.ipynb" - typer.echo(f"[colab] Saving notebook with outputs to '{output_file}'...") + info(f"[colab] Saving notebook with outputs to '{output_file}'...") with open(output_file, "w", encoding="utf-8") as f: nbformat.write(nb, f) + if failed: + raise typer.Exit(1) def repl( diff --git a/src/colab_cli/utils.py b/src/colab_cli/utils.py index 9943da0..b97f9e9 100644 --- a/src/colab_cli/utils.py +++ b/src/colab_cli/utils.py @@ -100,20 +100,25 @@ def render_display_data(data: dict) -> Union[Markdown, Text, None]: return None -def handle_image(image_b64: str, mime_type: str = "image/png", target_path: str = None): +def save_image( + image_b64: str, mime_type: str = "image/png", target_path: str = None +) -> str: + """Write a base64 image to `target_path` (or a temp file); return its path.""" image_bytes = base64.b64decode(image_b64) - # Print inline using Kitty protocol - print_kitty(image_bytes) - if target_path: - # If a target path is specified, save it there with open(target_path, "wb") as f: f.write(image_bytes) - print(f"\n[Image saved to: {target_path}]") - else: - # Save to temp file as fallback - ext = mime_type.split("/")[-1] - tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}") - tmp.write(image_bytes) - tmp.close() - print(f"\n[Image saved to: {tmp.name}]") + return target_path + # Save to temp file as fallback + ext = mime_type.split("/")[-1] + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}") + tmp.write(image_bytes) + tmp.close() + return tmp.name + + +def handle_image(image_b64: str, mime_type: str = "image/png", target_path: str = None): + # Print inline using Kitty protocol + print_kitty(base64.b64decode(image_b64)) + path = save_image(image_b64, mime_type, target_path) + print(f"\n[Image saved to: {path}]") diff --git a/tests/test_exec_json.py b/tests/test_exec_json.py new file mode 100644 index 0000000..03b115a --- /dev/null +++ b/tests/test_exec_json.py @@ -0,0 +1,252 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for `colab exec --json`, the machine-readable output mode.""" + +import base64 +import json +from unittest.mock import MagicMock + +import nbformat +import pytest +from nbformat.v4 import new_code_cell, new_notebook +from typer.testing import CliRunner + +from colab_cli.cli import app + +runner = CliRunner() + +SETUP_CODE = "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" + + +@pytest.fixture +def mock_runtime(mocker, mock_common_state): + session = MagicMock() + session.name = "s1" + session.url = "http://url" + session.token = "token" + session.kernel_id = None + session.session_id = None + mock_common_state.store.get.return_value = session + mock_common_state.resolve_session.return_value = "s1" + runtime_class = mocker.patch("colab_cli.commands.execution.ColabRuntime") + return runtime_class.return_value + + +def _script(runtime, per_code): + """Make `execute_code` stream and return the outputs mapped to each code.""" + + def execute_code(code, output_hook=None, **kwargs): + if code == SETUP_CODE: + return [] + outputs = per_code[code] + if isinstance(outputs, Exception): + raise outputs + if output_hook: + for out in outputs: + output_hook(out) + return outputs + + runtime.execute_code.side_effect = execute_code + + +def _records(stdout): + return [json.loads(line) for line in stdout.splitlines()] + + +def test_exec_json_success_emits_one_record(mock_runtime): + _script( + mock_runtime, + { + "run()": [ + {"output_type": "stream", "name": "stdout", "text": "hello\n"}, + {"output_type": "stream", "name": "stderr", "text": "warn\n"}, + { + "output_type": "execute_result", + "data": {"text/plain": "42"}, + "metadata": {}, + }, + ] + }, + ) + + result = runner.invoke(app, ["exec", "-s", "s1", "--json"], input="run()") + + assert result.exit_code == 0, result.output + assert _records(result.stdout) == [ + { + "cell_index": 0, + "cell_id": None, + "status": "ok", + "stdout": "hello\n", + "stderr": "warn\n", + "data": [{"output_type": "execute_result", "text/plain": "42"}], + "images": [], + "error": None, + } + ] + # The kernel's stderr is captured in the record, not replayed. + assert result.stderr == "" + + +def test_exec_json_error_strips_ansi_and_exits_nonzero(mock_runtime): + _script( + mock_runtime, + { + "boom()": [ + {"output_type": "stream", "name": "stdout", "text": "before\n"}, + { + "output_type": "error", + "ename": "ValueError", + "evalue": "bad", + "traceback": [ + "\x1b[0;31mValueError\x1b[0m Traceback", + "\x1b[1;32m----> 1\x1b[0m boom()", + ], + }, + ] + }, + ) + + result = runner.invoke(app, ["exec", "-s", "s1", "--json"], input="boom()") + + assert result.exit_code == 1 + (record,) = _records(result.stdout) + assert record["status"] == "error" + assert record["stdout"] == "before\n" + assert record["error"] == { + "ename": "ValueError", + "evalue": "bad", + "traceback": "ValueError Traceback\n----> 1 boom()", + } + assert "\x1b" not in result.stdout + assert result.stderr == "" + + +def test_exec_json_saves_images_and_reports_paths(mock_runtime, tmp_path): + png = base64.b64encode(b"png-bytes").decode("ascii") + _script( + mock_runtime, + { + "plot()": [ + { + "output_type": "display_data", + "data": {"image/png": png, "text/plain": "
"}, + "metadata": {}, + } + ] + }, + ) + target = tmp_path / "plot.png" + + result = runner.invoke( + app, + ["exec", "-s", "s1", "--json", "--output-image", str(target)], + input="plot()", + ) + + assert result.exit_code == 0, result.output + (record,) = _records(result.stdout) + assert record["images"] == [str(target)] + assert record["data"] == [{"output_type": "display_data", "text/plain": "
"}] + assert png not in result.stdout + assert target.read_bytes() == b"png-bytes" + + +def test_exec_json_notebook_emits_record_per_cell(mock_runtime, tmp_path): + nb = new_notebook( + cells=[ + new_code_cell("ok()", id="cell-a"), + new_code_cell("fail()", id="cell-b"), + new_code_cell("after()", id="cell-c"), + ] + ) + path = tmp_path / "nb.ipynb" + nbformat.write(nb, str(path)) + _script( + mock_runtime, + { + "ok()": [{"output_type": "stream", "name": "stdout", "text": "1\n"}], + "fail()": [ + { + "output_type": "error", + "ename": "NameError", + "evalue": "fail", + "traceback": [], + } + ], + "after()": [], + }, + ) + + result = runner.invoke(app, ["exec", "-s", "s1", "--json", "-f", str(path)]) + + assert result.exit_code == 1 + records = _records(result.stdout) + assert [(r["cell_index"], r["cell_id"], r["status"]) for r in records] == [ + (0, "cell-a", "ok"), + (1, "cell-b", "error"), + (2, "cell-c", "ok"), + ] + # Progress chatter moves to stderr so stdout stays pure NDJSON. + assert "[colab] Executing cell 1/3" in result.stderr + assert "_output.ipynb" in result.stderr + + +def test_exec_json_reports_client_side_failure(mock_runtime): + _script(mock_runtime, {"slow()": TimeoutError("no output for 30s")}) + + result = runner.invoke(app, ["exec", "-s", "s1", "--json"], input="slow()") + + assert result.exit_code == 1 + (record,) = _records(result.stdout) + assert record["status"] == "error" + assert record["error"] == { + "ename": "TimeoutError", + "evalue": "no output for 30s", + "traceback": "", + } + + +def test_exec_json_lost_session_keeps_stdout_clean(mock_runtime, mock_common_state): + mock_runtime.execute_code.side_effect = Exception("404 Not Found") + + result = runner.invoke(app, ["exec", "-s", "s1", "--json"], input="print(1)") + + assert result.exit_code == 1 + assert result.stdout == "" + assert "appears to be lost" in result.stderr + mock_common_state.prune_session.assert_called_once_with("s1") + + +def test_exec_without_json_still_exits_zero_on_kernel_error(mock_runtime): + """`--json` is opt-in: the default exit code and streams are unchanged.""" + _script( + mock_runtime, + { + "boom()": [ + { + "output_type": "error", + "ename": "ValueError", + "evalue": "bad", + "traceback": [], + } + ] + }, + ) + + result = runner.invoke(app, ["exec", "-s", "s1"], input="boom()") + + assert result.exit_code == 0 + assert "ValueError: bad" in result.stderr