From de037d5021ca6ddad8060fc0e4060b95f2ca4133 Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 16 Sep 2026 10:47:36 -0700 Subject: [PATCH 1/7] feat(pair): add a client for live notebook servers Stream scratchpad execution over SSE, list sessions, resolve a session from a notebook path, and read server URLs from the local registry. Tokens travel in headers only, and every server response is validated before use. --- marimo/_cli/pair/client.py | 393 ++++++++++ tests/_cli/fixtures/pair/execute-failure.sse | 6 + tests/_cli/fixtures/pair/execute-success.sse | 6 + tests/_cli/test_pair_client.py | 748 +++++++++++++++++++ 4 files changed, 1153 insertions(+) create mode 100644 marimo/_cli/pair/client.py create mode 100644 tests/_cli/fixtures/pair/execute-failure.sse create mode 100644 tests/_cli/fixtures/pair/execute-success.sse create mode 100644 tests/_cli/test_pair_client.py diff --git a/marimo/_cli/pair/client.py b/marimo/_cli/pair/client.py new file mode 100644 index 00000000000..5105ff3ddb3 --- /dev/null +++ b/marimo/_cli/pair/client.py @@ -0,0 +1,393 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +import http.client +import json +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING +from urllib.parse import urlsplit + +from marimo._server.api.utils import format_url_host +from marimo._server.server_registry import _servers_dir + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator, Mapping + from http.client import HTTPResponse + from pathlib import Path + from typing import TextIO + + +class PairError(Exception): + """A failure that the CLI reports on stderr.""" + + +class PairInputError(PairError): + """Invalid local input that prevents a server operation.""" + + +class NoSessionError(PairError): + def __init__(self, message: str, *, url: str) -> None: + super().__init__(message) + self.url = url + + +class AmbiguousSessionError(PairError): + def __init__( + self, message: str, *, url: str, candidates: tuple[str, ...] + ) -> None: + super().__init__(message) + self.url = url + self.candidates = candidates + + +class StaleSessionError(PairError): + """The server rejected the session ID as unknown.""" + + +@dataclass(frozen=True) +class SSEEvent: + name: str + data: str + + +@dataclass(frozen=True) +class ExecutionResult: + success: bool + output: dict[str, str] | None + stdout: str + stderr: str + + +def load_token( + token_file: Path | None, environ: Mapping[str, str] +) -> str | None: + if token_file is None: + return environ.get("MARIMO_TOKEN") or None + + try: + token = token_file.read_text(encoding="utf-8").rstrip("\r\n") + except (OSError, UnicodeError) as error: + raise PairInputError("Could not read the token file.") from error + if not token: + raise PairInputError("The token file is empty.") + return token + + +def display_url(url: str) -> str: + """Return a URL without credentials, query values, or a fragment.""" + try: + parsed = urlsplit(url) + except ValueError: + return "" + return parsed._replace( + netloc=parsed.netloc.rsplit("@", 1)[-1], query="", fragment="" + ).geturl() + + +def _endpoint_url(url: str, path: str) -> str: + try: + parsed = urlsplit(url) + except ValueError as error: + raise PairInputError("The server URL is invalid.") from error + return parsed._replace( + path=f"{parsed.path.rstrip('/')}{path}", fragment="" + ).geturl() + + +def iter_sse(lines: Iterable[bytes]) -> Iterator[SSEEvent]: + name = "message" + data: list[str] = [] + + for raw_line in lines: + line = raw_line.decode("utf-8").rstrip("\r\n") + if not line: + if data: + yield SSEEvent(name=name, data="\n".join(data)) + name = "message" + data = [] + continue + if line.startswith(":"): + continue + + field, separator, value = line.partition(":") + if separator and value.startswith(" "): + value = value[1:] + if field == "event": + name = value + elif field == "data": + data.append(value) + + if data: + yield SSEEvent(name=name, data="\n".join(data)) + + +def open_response( + *, method: str, url: str, headers: dict[str, str], body: bytes | None +) -> HTTPResponse: + try: + parsed = urlsplit(url) + port = parsed.port + except ValueError as error: + raise PairInputError("The server URL is invalid.") from error + if parsed.scheme not in ("http", "https") or parsed.hostname is None: + raise PairInputError("The server URL must use http or https.") + if port is not None and not 1 <= port <= 65535: + raise PairInputError("The server URL is invalid.") + + connection_type = ( + http.client.HTTPSConnection + if parsed.scheme == "https" + else http.client.HTTPConnection + ) + try: + connection = connection_type( + parsed.hostname, + port, + timeout=5.0, + ) + connection.connect() + if connection.sock is not None: + connection.sock.settimeout(None) + path = parsed.path or "/" + if parsed.query: + path = f"{path}?{parsed.query}" + connection.request(method, path, body=body, headers=headers) + response = connection.getresponse() + except (ValueError, http.client.InvalidURL) as error: + raise PairInputError("The server URL is invalid.") from error + except (OSError, http.client.HTTPException) as error: + raise PairError("Could not connect to the server.") from error + + _raise_for_status(response) + return response + + +def _response_detail(response: HTTPResponse) -> str | None: + try: + payload = json.loads(response.read()) + except (OSError, UnicodeError, ValueError, http.client.HTTPException): + return None + if isinstance(payload, dict): + detail = payload.get("detail") + if isinstance(detail, str): + return detail + return None + + +def _raise_for_status(response: HTTPResponse) -> None: + if response.status in (401, 403): + response.close() + raise PairError("Authentication failed.") + if 200 <= response.status < 300: + return + detail = _response_detail(response) + response.close() + if isinstance(detail, str) and detail.startswith("Invalid session id"): + raise StaleSessionError(detail) + if detail == "Missing Marimo-Session-Id header": + raise PairError("Internal: should not happen after resolution.") + if detail: + raise PairError(detail) + raise PairError(f"Server returned {response.status}.") + + +def execute( + *, + url: str, + session_id: str, + token: str | None, + code: str, + stdout: TextIO, + stderr: TextIO, + stream: bool, +) -> ExecutionResult: + request_url = _endpoint_url(url, "/api/kernel/execute") + headers = { + "Content-Type": "application/json", + "Marimo-Session-Id": session_id, + } + if token is not None: + headers["Authorization"] = f"Bearer {token}" + body = json.dumps({"code": code}).encode("utf-8") + response = open_response( + method="POST", + url=request_url, + headers=headers, + body=body, + ) + stdout_parts: list[str] = [] + stderr_parts: list[str] = [] + + def write_event(target: TextIO, parts: list[str], value: str) -> None: + if stream: + target.write(value) + target.flush() + else: + parts.append(value) + + try: + try: + for event in iter_sse(response): + payload = json.loads(event.data) + if event.name == "stdout": + write_event(stdout, stdout_parts, payload["data"]) + elif event.name == "stderr": + write_event(stderr, stderr_parts, payload["data"]) + elif event.name == "done": + raw_output = payload.get("output") + data = "" + mimetype = "text/plain" + if isinstance(raw_output, dict): + data = str(raw_output.get("data") or "") + mimetype = str( + raw_output.get("mimetype") or "text/plain" + ) + output = ( + None + if not data + else {"mimetype": mimetype, "data": data} + ) + if stream and output is not None: + stdout.write(f"{output['data']}\n") + return ExecutionResult( + success=bool(payload["success"]), + output=output, + stdout="".join(stdout_parts), + stderr="".join(stderr_parts), + ) + except (OSError, http.client.HTTPException) as error: + raise PairError( + "The execution response ended before completion was confirmed." + ) from error + + raise PairError( + "The execution response ended before completion was confirmed." + ) + finally: + response.close() + + +def list_sessions( + *, url: str, token: str | None +) -> dict[str, dict[str, str | None]]: + request_url = _endpoint_url(url, "/api/sessions") + headers = {} + if token is not None: + headers["Authorization"] = f"Bearer {token}" + response = open_response( + method="GET", + url=request_url, + headers=headers, + body=None, + ) + try: + try: + payload = json.load(response) + except ( + OSError, + UnicodeError, + json.JSONDecodeError, + http.client.HTTPException, + ) as error: + raise PairError( + f"Unexpected response from {display_url(request_url)}." + ) from error + finally: + response.close() + + message = f"Unexpected response from {display_url(request_url)}." + if not isinstance(payload, dict): + raise PairError(message) + + sessions: dict[str, dict[str, str | None]] = {} + for session_id, session in payload.items(): + if not isinstance(session_id, str) or not isinstance(session, dict): + raise PairError(message) + if "filename" not in session or "path" not in session: + raise PairError(message) + filename = session["filename"] + path = session["path"] + if (filename is not None and not isinstance(filename, str)) or ( + path is not None and not isinstance(path, str) + ): + raise PairError(message) + sessions[session_id] = {"filename": filename, "path": path} + return sessions + + +def resolve_session(*, url: str, token: str | None, file: str | None) -> str: + sessions = list_sessions(url=url, token=token) + safe_url = display_url(url) + if file is None: + candidates = list(sessions) + else: + for field, value in ( + ("path", file), + ("filename", file), + ("path", os.path.abspath(file)), + ): + candidates = [ + session_id + for session_id, session in sessions.items() + if session[field] == value + ] + if candidates: + break + + if len(candidates) == 1: + return candidates[0] + if not candidates: + if file is None: + message = f"No running session on {safe_url}." + else: + message = ( + f"No running session for notebook '{file}' on {safe_url}." + ) + raise NoSessionError(message, url=safe_url) + + candidates.sort() + if file is None: + message = f"Server {safe_url} has {len(candidates)} running sessions." + else: + message = ( + f"Notebook '{file}' has {len(candidates)} running sessions " + f"on {safe_url}." + ) + raise AmbiguousSessionError( + message, url=safe_url, candidates=tuple(candidates) + ) + + +def registry_urls() -> list[str]: + import psutil + + urls: list[str] = [] + for path in sorted(_servers_dir().glob("*.json")): + try: + with path.open(encoding="utf-8") as file: + entry = json.load(file) + except (OSError, UnicodeError, json.JSONDecodeError): + continue + if not isinstance(entry, dict): + continue + pid = entry.get("pid") + host = entry.get("host") + port = entry.get("port") + base_url = entry.get("base_url") + if ( + type(pid) is not int + or not psutil.pid_exists(pid) + or not isinstance(host, str) + or type(port) is not int + or not 1 <= port <= 65535 + or not isinstance(base_url, str) + ): + continue + + url_host = format_url_host(host, port, route_bind_all_to_loopback=True) + if port == 80: + urls.append(f"http://{url_host}{base_url}") + else: + urls.append(f"http://{url_host}:{port}{base_url}") + return urls diff --git a/tests/_cli/fixtures/pair/execute-failure.sse b/tests/_cli/fixtures/pair/execute-failure.sse new file mode 100644 index 00000000000..60e96bb3525 --- /dev/null +++ b/tests/_cli/fixtures/pair/execute-failure.sse @@ -0,0 +1,6 @@ +event: stderr +data: {"data":"ValueError: boom\n"} + +event: done +data: {"success":false,"output":{"mimetype":"text/plain","data":""}} + diff --git a/tests/_cli/fixtures/pair/execute-success.sse b/tests/_cli/fixtures/pair/execute-success.sse new file mode 100644 index 00000000000..f01bc8076e9 --- /dev/null +++ b/tests/_cli/fixtures/pair/execute-success.sse @@ -0,0 +1,6 @@ +event: stdout +data: {"data":"hello\n"} + +event: done +data: {"success":true,"output":{"mimetype":"text/plain","data":"2"}} + diff --git a/tests/_cli/test_pair_client.py b/tests/_cli/test_pair_client.py new file mode 100644 index 00000000000..a6f057a3ab3 --- /dev/null +++ b/tests/_cli/test_pair_client.py @@ -0,0 +1,748 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +import http.client +import io +import json +import os +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import pytest + +from marimo._cli.pair import client +from marimo._cli.pair.client import ( + AmbiguousSessionError, + NoSessionError, + PairError, + PairInputError, + SSEEvent, + StaleSessionError, +) + +if TYPE_CHECKING: + from collections.abc import Iterator + + +FIXTURES = Path(__file__).parent / "fixtures" / "pair" + + +class RecordingStream(io.StringIO): + def __init__(self, name: str, records: list[tuple[str, str]]) -> None: + super().__init__() + self.name = name + self.records = records + + def write(self, value: str) -> int: + self.records.append((self.name, value)) + return super().write(value) + + def flush(self) -> None: + self.records.append((self.name, "flush")) + super().flush() + + +class RaisingResponse: + def __init__(self, lines: list[bytes], error: BaseException) -> None: + self.lines = lines + self.error = error + self.closed = False + + def __iter__(self) -> Iterator[bytes]: + yield from self.lines + raise self.error + + def close(self) -> None: + self.closed = True + + +def _fixture(name: str) -> bytes: + return (FIXTURES / name).read_bytes() + + +def _patch_response( + monkeypatch: pytest.MonkeyPatch, response: Any +) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + + def fake_open_response(**kwargs: Any) -> Any: + calls.append(kwargs) + return response + + monkeypatch.setattr(client, "open_response", fake_open_response) + return calls + + +def test_load_token_without_file_or_environment() -> None: + assert client.load_token(None, {}) is None + + +def test_load_token_from_environment() -> None: + assert client.load_token(None, {"MARIMO_TOKEN": "environment"}) == ( + "environment" + ) + + +def test_load_token_file_wins_and_trims_newline(tmp_path: Path) -> None: + token_file = tmp_path / "token.txt" + token_file.write_text(" file token \r\n", encoding="utf-8") + + assert ( + client.load_token(token_file, {"MARIMO_TOKEN": "environment"}) + == " file token " + ) + + +def test_load_token_rejects_unreadable_file(tmp_path: Path) -> None: + with pytest.raises(PairInputError, match="Could not read the token file"): + client.load_token(tmp_path, {"MARIMO_TOKEN": "environment"}) + + +def test_load_token_rejects_empty_file(tmp_path: Path) -> None: + token_file = tmp_path / "token.txt" + token_file.write_text("\n", encoding="utf-8") + + with pytest.raises(PairInputError, match="The token file is empty"): + client.load_token(token_file, {}) + + +def test_display_url_removes_credentials_query_and_fragment() -> None: + assert ( + client.display_url( + "https://user:password@example.com/base?access_token=secret#part" + ) + == "https://example.com/base" + ) + assert client.display_url("http://[::1") == "" + + +@pytest.mark.parametrize( + ("fixture", "expected"), + [ + ( + "execute-success.sse", + [ + SSEEvent("stdout", '{"data":"hello\\n"}'), + SSEEvent( + "done", + '{"success":true,"output":{"mimetype":"text/plain","data":"2"}}', + ), + ], + ), + ( + "execute-failure.sse", + [ + SSEEvent("stderr", '{"data":"ValueError: boom\\n"}'), + SSEEvent( + "done", + '{"success":false,"output":{"mimetype":"text/plain","data":""}}', + ), + ], + ), + ], +) +def test_iter_sse_fixtures(fixture: str, expected: list[SSEEvent]) -> None: + assert list(client.iter_sse(io.BytesIO(_fixture(fixture)))) == expected + + +def test_iter_sse_handles_crlf_comments_and_multiline_data() -> None: + lines = [ + b": keep-alive\r\n", + b"event: custom\r\n", + b"data: first\r\n", + b"data: second\r\n", + b"\r\n", + ] + + assert list(client.iter_sse(lines)) == [ + SSEEvent("custom", "first\nsecond") + ] + + +def test_iter_sse_dispatches_unterminated_final_record() -> None: + assert list(client.iter_sse([b"data: final"])) == [ + SSEEvent("message", "final") + ] + + +def test_execute_sends_request_and_streams_in_event_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = io.BytesIO( + b"event: stdout\n" + b'data: {"data":"out"}\n\n' + b"event: stderr\n" + b'data: {"data":"err"}\n\n' + b"event: done\n" + b'data: {"success":true,"output":{"data":"result"}}\n\n' + ) + calls = _patch_response(monkeypatch, response) + records: list[tuple[str, str]] = [] + stdout = RecordingStream("stdout", records) + stderr = RecordingStream("stderr", records) + + result = client.execute( + url="https://example.com/base/?access_token=query-token#part", + session_id="session-1", + token="secret-token", + code="print(1)", + stdout=stdout, + stderr=stderr, + stream=True, + ) + + assert result == client.ExecutionResult( + success=True, + output={"mimetype": "text/plain", "data": "result"}, + stdout="", + stderr="", + ) + assert records == [ + ("stdout", "out"), + ("stdout", "flush"), + ("stderr", "err"), + ("stderr", "flush"), + ("stdout", "result\n"), + ] + assert calls == [ + { + "method": "POST", + "url": "https://example.com/base/api/kernel/execute?access_token=query-token", + "headers": { + "Content-Type": "application/json", + "Marimo-Session-Id": "session-1", + "Authorization": "Bearer secret-token", + }, + "body": json.dumps({"code": "print(1)"}).encode(), + } + ] + assert "secret-token" not in calls[0]["url"] + assert b"secret-token" not in calls[0]["body"] + assert response.closed + + +def test_execute_omits_authorization_without_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = io.BytesIO(_fixture("execute-success.sse")) + calls = _patch_response(monkeypatch, response) + + client.execute( + url="http://localhost:2718", + session_id="session-1", + token=None, + code="1 + 1", + stdout=io.StringIO(), + stderr=io.StringIO(), + stream=True, + ) + + assert len(calls) == 1 + assert "Authorization" not in calls[0]["headers"] + + +@pytest.mark.parametrize( + "url", + [ + "ftp://localhost/not-http", + "http:///missing-host", + "http://localhost:0", + "http://localhost:not-a-port?access_token=secret", + "http://[::1", + ], +) +def test_open_response_rejects_invalid_url_without_exposing_it( + url: str, +) -> None: + with pytest.raises(PairInputError) as exc_info: + client.open_response( + method="GET", + url=url, + headers={}, + body=None, + ) + + assert str(exc_info.value) in ( + "The server URL is invalid.", + "The server URL must use http or https.", + ) + assert "secret" not in str(exc_info.value) + + +def test_execute_buffers_output_until_done( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stdout = io.StringIO() + stderr = io.StringIO() + + class InspectingResponse(io.BytesIO): + def readline(self, size: int = -1) -> bytes: + line = super().readline(size) + if line == b"event: done\n": + assert stdout.getvalue() == "" + assert stderr.getvalue() == "" + return line + + response = InspectingResponse(_fixture("execute-success.sse")) + calls = _patch_response(monkeypatch, response) + + result = client.execute( + url="http://localhost:2718", + session_id="session-1", + token=None, + code="1 + 1", + stdout=stdout, + stderr=stderr, + stream=False, + ) + + assert result.success + assert result.stdout == "hello\n" + assert result.stderr == "" + assert result.output == {"mimetype": "text/plain", "data": "2"} + assert stdout.getvalue() == "" + assert stderr.getvalue() == "" + assert len(calls) == 1 + + +def test_execute_returns_unsuccessful_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = io.BytesIO(_fixture("execute-failure.sse")) + calls = _patch_response(monkeypatch, response) + stderr = io.StringIO() + + result = client.execute( + url="http://localhost:2718", + session_id="session-1", + token=None, + code="raise ValueError", + stdout=io.StringIO(), + stderr=stderr, + stream=False, + ) + + assert result == client.ExecutionResult( + success=False, + output=None, + stdout="", + stderr="ValueError: boom\n", + ) + assert stderr.getvalue() == "" + assert len(calls) == 1 + + +def test_execute_rejects_missing_done( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = io.BytesIO(b'event: stdout\ndata: {"data":"partial"}\n\n') + calls = _patch_response(monkeypatch, response) + + with pytest.raises(PairError, match="ended before completion"): + client.execute( + url="http://localhost:2718", + session_id="session-1", + token=None, + code="print(1)", + stdout=io.StringIO(), + stderr=io.StringIO(), + stream=True, + ) + + assert len(calls) == 1 + + +def test_execute_keeps_buffered_output_after_read_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = RaisingResponse( + [b"event: stdout\n", b'data: {"data":"partial"}\n', b"\n"], + OSError("connection lost"), + ) + calls = _patch_response(monkeypatch, response) + stdout = io.StringIO() + + with pytest.raises(PairError, match="ended before completion"): + client.execute( + url="http://localhost:2718", + session_id="session-1", + token=None, + code="print(1)", + stdout=stdout, + stderr=io.StringIO(), + stream=True, + ) + + assert stdout.getvalue() == "partial" + assert response.closed + assert len(calls) == 1 + + +def test_execute_rejects_truncated_http_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = RaisingResponse([], http.client.IncompleteRead(b"partial")) + _patch_response(monkeypatch, response) + + with pytest.raises(PairError, match="ended before completion"): + client.execute( + url="http://localhost:2718", + session_id="session-1", + token=None, + code="print(1)", + stdout=io.StringIO(), + stderr=io.StringIO(), + stream=True, + ) + + assert response.closed + + +def test_execute_closes_response_after_keyboard_interrupt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = RaisingResponse([], KeyboardInterrupt()) + calls = _patch_response(monkeypatch, response) + + with pytest.raises(KeyboardInterrupt): + client.execute( + url="http://localhost:2718", + session_id="session-1", + token=None, + code="while True: pass", + stdout=io.StringIO(), + stderr=io.StringIO(), + stream=True, + ) + + assert response.closed + assert len(calls) == 1 + + +def test_list_sessions_sends_request_and_parses_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = io.BytesIO( + json.dumps( + { + "session-1": { + "filename": "analysis.py", + "path": "/work/analysis.py", + "state": "idle", + }, + "session-2": {"filename": None, "path": None}, + } + ).encode() + ) + calls = _patch_response(monkeypatch, response) + + sessions = client.list_sessions( + url=( + "https://example.com/base/" + "?access_token=query-token#ignored-fragment" + ), + token="secret-token", + ) + + assert sessions == { + "session-1": { + "filename": "analysis.py", + "path": "/work/analysis.py", + }, + "session-2": {"filename": None, "path": None}, + } + assert calls == [ + { + "method": "GET", + "url": ( + "https://example.com/base/api/sessions" + "?access_token=query-token" + ), + "headers": {"Authorization": "Bearer secret-token"}, + "body": None, + } + ] + assert response.closed + + +def test_list_sessions_omits_authorization_without_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = io.BytesIO(b"{}") + calls = _patch_response(monkeypatch, response) + + assert client.list_sessions(url="http://localhost:2718", token=None) == {} + + assert calls[0]["headers"] == {} + assert response.closed + + +@pytest.mark.parametrize( + "payload", + [ + [], + {"session-1": None}, + {"session-1": {}}, + {"session-1": {"filename": "analysis.py"}}, + {"session-1": {"path": "/work/analysis.py"}}, + {"session-1": {"filename": [], "path": None}}, + {"session-1": {"filename": None, "path": []}}, + ], +) +def test_list_sessions_rejects_malformed_response( + monkeypatch: pytest.MonkeyPatch, + payload: object, +) -> None: + response = io.BytesIO(json.dumps(payload).encode()) + _patch_response(monkeypatch, response) + + with pytest.raises(PairError) as exc_info: + client.list_sessions( + url="http://user:password@localhost:2718/" + "?access_token=secret#fragment", + token=None, + ) + + assert str(exc_info.value) == ( + "Unexpected response from http://localhost:2718/api/sessions." + ) + assert "password" not in str(exc_info.value) + assert "secret" not in str(exc_info.value) + assert response.closed + + +def test_list_sessions_rejects_invalid_json( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = io.BytesIO(b"{") + _patch_response(monkeypatch, response) + + with pytest.raises(PairError) as exc_info: + client.list_sessions( + url="http://localhost:2718/?access_token=secret", token=None + ) + + assert str(exc_info.value) == ( + "Unexpected response from http://localhost:2718/api/sessions." + ) + assert "secret" not in str(exc_info.value) + assert response.closed + + +def test_registry_urls_skips_invalid_entries( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + (tmp_path / "valid.json").write_text( + json.dumps( + { + "pid": os.getpid(), + "host": "127.0.0.1", + "port": 2718, + "base_url": "", + } + ), + encoding="utf-8", + ) + (tmp_path / "stale.json").write_text( + json.dumps( + { + "pid": 2**31 - 1, + "host": "127.0.0.1", + "port": 2720, + "base_url": "", + } + ), + encoding="utf-8", + ) + (tmp_path / "invalid.json").write_text("{", encoding="utf-8") + (tmp_path / "missing.json").write_text( + json.dumps({"port": 2719}), encoding="utf-8" + ) + (tmp_path / "bad-port.json").write_text( + json.dumps( + { + "pid": os.getpid(), + "host": "127.0.0.1", + "port": 70000, + "base_url": "", + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(client, "_servers_dir", lambda: tmp_path) + + assert client.registry_urls() == ["http://localhost:2718"] + + +def test_registry_urls_formats_prefix_and_standard_ports( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + entries = [ + ("http.json", "0.0.0.0", 80, "/prefix"), + ("https.json", "::", 443, ""), + ] + for filename, host, port, base_url in entries: + (tmp_path / filename).write_text( + json.dumps( + { + "pid": os.getpid(), + "host": host, + "port": port, + "base_url": base_url, + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(client, "_servers_dir", lambda: tmp_path) + + assert client.registry_urls() == [ + "http://localhost/prefix", + "http://localhost:443", + ] + + +class FakeStatusResponse: + def __init__(self, status: int, body: bytes = b"") -> None: + self.status = status + self.body = body + self.closed = False + + def read(self) -> bytes: + return self.body + + def close(self) -> None: + self.closed = True + + +def test_raise_for_status_maps_invalid_session_id() -> None: + response = FakeStatusResponse( + 500, json.dumps({"detail": "Invalid session id: s_old"}).encode() + ) + + with pytest.raises(StaleSessionError, match="Invalid session id"): + client._raise_for_status(response) + + assert response.closed + + +def test_raise_for_status_uses_json_detail() -> None: + response = FakeStatusResponse( + 500, + json.dumps({"detail": "Missing Marimo-Session-Id header"}).encode(), + ) + + with pytest.raises( + PairError, match="Internal: should not happen after resolution" + ): + client._raise_for_status(response) + + +def test_resolve_session_matches_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + client, + "list_sessions", + lambda **_kwargs: { + "s_one": {"filename": "analysis.py", "path": "/work/analysis.py"} + }, + ) + + assert ( + client.resolve_session( + url="http://one", token=None, file="/work/analysis.py" + ) + == "s_one" + ) + + +def test_resolve_session_matches_filename( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + client, + "list_sessions", + lambda **_kwargs: { + "s_one": {"filename": "nb.py", "path": "/work/nb.py"} + }, + ) + + assert ( + client.resolve_session(url="http://one", token=None, file="nb.py") + == "s_one" + ) + + +def test_resolve_session_matches_abspath( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + notebook = tmp_path / "nb.py" + monkeypatch.chdir(tmp_path) + monkeypatch.setattr( + client, + "list_sessions", + lambda **_kwargs: { + "s_one": { + "filename": "other.py", + "path": str(notebook.resolve()), + } + }, + ) + + assert ( + client.resolve_session(url="http://one", token=None, file="nb.py") + == "s_one" + ) + + +def test_resolve_session_zero_matches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(client, "list_sessions", lambda **_kwargs: {}) + + with pytest.raises(NoSessionError, match=r"gone\.py") as exc_info: + client.resolve_session( + url="http://user:password@one?access_token=secret", + token=None, + file="gone.py", + ) + + assert exc_info.value.url == "http://one" + assert "password" not in str(exc_info.value) + assert "secret" not in str(exc_info.value) + + +def test_resolve_session_stops_at_first_matching_tier( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + client, + "list_sessions", + lambda **_kwargs: { + "s_exact": {"filename": "other.py", "path": "nb.py"}, + "s_filename": {"filename": "nb.py", "path": "/work/nb.py"}, + }, + ) + + assert ( + client.resolve_session(url="http://one", token=None, file="nb.py") + == "s_exact" + ) + + +def test_resolve_session_many_matches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + client, + "list_sessions", + lambda **_kwargs: { + "s_b": {"filename": "nb.py", "path": "/b/nb.py"}, + "s_a": {"filename": "nb.py", "path": "/a/nb.py"}, + }, + ) + + with pytest.raises(AmbiguousSessionError) as exc_info: + client.resolve_session(url="http://one", token=None, file="nb.py") + + assert exc_info.value.candidates == ("s_a", "s_b") From 22be2ebbf52b6bb5dc5d11205958cf4ba0607e13 Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 16 Sep 2026 10:47:48 -0700 Subject: [PATCH 2/7] feat(pair): add execute, docs, and notebook list commands execute runs Python in a live kernel and prints one JSON result, or one JSON failure with a next step an agent can run. notebook list finds servers and sessions. --help is a copyable procedure. prompt no longer checks for an installed skill, because the CLI replaces it. --- marimo/_cli/pair/commands.py | 702 ++++++++++---- tests/_cli/test_cli_pair.py | 1672 ++++++++++++++++++++++++++-------- 2 files changed, 1828 insertions(+), 546 deletions(-) diff --git a/marimo/_cli/pair/commands.py b/marimo/_cli/pair/commands.py index 122edcef904..99694607141 100644 --- a/marimo/_cli/pair/commands.py +++ b/marimo/_cli/pair/commands.py @@ -2,17 +2,33 @@ from __future__ import annotations import hashlib +import json import os -import shlex -from dataclasses import dataclass, field +import re +import sys from pathlib import Path import click from marimo._cli.help_formatter import ColoredCommand, ColoredGroup +from marimo._cli.pair.client import ( + AmbiguousSessionError, + NoSessionError, + PairError, + PairInputError, + StaleSessionError, + display_url, + execute as execute_code, + list_sessions, + load_token, + registry_urls, + resolve_session, +) +from marimo._server.ai.skills import utils as skills_utils -SKILL_NAME = "marimo-pair" -SKILL_FILE = "SKILL.md" +_REFERENCES_DIR = ( + Path(skills_utils.__file__).parent / "marimo-pair" / "references" +) _cached_token_dir: Path | None = None @@ -27,147 +43,385 @@ def _token_dir() -> Path: return _cached_token_dir -@dataclass(frozen=True) -class AgentConfig: - name: str - skill_dirs: list[Path] = field(default_factory=list) +def _doc_topics() -> dict[str, str]: + return { + path.stem: path.read_text(encoding="utf-8") + .splitlines()[0] + .removeprefix("# ") + for path in sorted(_REFERENCES_DIR.glob("*.md")) + } - def has_skill(self) -> bool: - for directory in self.skill_dirs: - try: - if (directory / SKILL_NAME / SKILL_FILE).exists(): - return True - except OSError: - # Skill detection is advisory. An inaccessible directory must - # not prevent marimo from generating the pairing prompt. - continue - return False +class _DocsCommand(ColoredCommand): + def format_epilog( + self, ctx: click.Context, formatter: click.HelpFormatter + ) -> None: + del ctx + formatter.write_paragraph() + formatter.write_text("Available topics:") + with formatter.indentation(): + formatter.write_dl(list(_doc_topics().items())) -def _claude_skill_dirs() -> list[Path]: - """Return all directories where a Claude Code skill may be installed. - Skills can be installed directly or bundled in a marketplace plugin in - both the global (`~/.claude`) and local (`.claude`) config directories. - """ - roots = [Path.home() / ".claude", Path.cwd() / ".claude"] - subdirs = ["skills", "plugins", str(Path("plugins") / "marketplaces")] - return [ - *[root / sub for root in roots for sub in subdirs], - *[ - skill_dir - for root in roots - for skill_dir in _plugin_skill_dirs(root) - ], - ] +@click.group( + cls=ColoredGroup, + help="""Pair with a live marimo notebook. + + \b + Workflow: + If you do not have the server URL and session id: + marimo pair notebook list + marimo pair execute --url --session --code-file - <<'PY' + import marimo._code_mode as cm + async with cm.get_context() as ctx: + cid = ctx.create_cell("x = 1") + ctx.run_cell(cid) + PY + marimo pair execute --url --session --code-file - <<'PY' + import marimo._code_mode as cm + async with cm.get_context() as ctx: + cell = ctx.cells[""] + print(cell.status, cell.errors, [o.data for o in cell.console_outputs]) + PY + + \b + Rules: + Cells are the unit of work. The scratchpad is temporary; only cm edits persist. + Cells do not run on creation. Call run_cell after create_cell or edit_cell. + Use async with. Do not await ctx methods. + Session IDs change when the page reloads. If execute reports a stale + session, run notebook list again. + + \b + Code-mode API (this marimo version): + ctx.cells # each has .id .code .status .errors .console_outputs + ctx.create_cell(code) # returns the new cell id + ctx.edit_cell(cid, code) + ctx.run_cell(cid) + ctx.delete_cell(cid) + If a cm call fails, run help(cm): + marimo pair execute --url --session -c 'import marimo._code_mode as cm; help(cm)' + """, +) +def pair() -> None: + pass -def _plugin_skill_dirs(root: Path) -> list[Path]: - """Return skill directories from marketplace and cached plugins.""" - plugins = root / "plugins" - return [ - *plugins.glob("marketplaces/*/skills"), - *plugins.glob(f"cache/*/{SKILL_NAME}/*/skills"), - ] +@click.command( + cls=ColoredCommand, + help="""Run Python in the selected live notebook kernel's scratchpad.""", + short_help="""Run Python in a live notebook session.""", +) +@click.option( + "--url", + required=True, + metavar="URL", + help="Server URL.", +) +@click.option( + "--session", + "session_id", + required=False, + metavar="ID", + help="Current session ID. Resolved from --file when omitted.", +) +@click.option( + "--file", + "file_path", + required=False, + metavar="PATH", + help="Notebook path or file key. Used to resolve --session when omitted.", +) +@click.option( + "--token-file", + type=click.Path(path_type=Path, dir_okay=False), + metavar="PATH", + help="Read the server token from a local file. Otherwise use MARIMO_TOKEN, if set.", +) +@click.option( + "-c", + "code", + help="Inline Python.", +) +@click.option( + "--code-file", + metavar="PATH", + help="Read Python from a UTF-8 file, or from stdin when PATH is '-'. Supply exactly one input option.", +) +@click.option( + "--stream", + "stream", + is_flag=True, + help="Write stdout and stderr as they arrive. Default: print one JSON result.", +) +@click.pass_context +def execute( + ctx: click.Context, + url: str, + session_id: str | None, + file_path: str | None, + token_file: Path | None, + code: str | None, + code_file: str | None, + stream: bool, +) -> None: + if (code is None) == (code_file is None): + raise click.UsageError("Specify -c or --code-file.") + if code_file == "-": + code = sys.stdin.read() + elif code_file is not None: + try: + code = Path(code_file).read_text(encoding="utf-8") + except (OSError, UnicodeError) as error: + raise click.UsageError("Could not read the code file.") from error + assert code is not None + if not code: + raise click.UsageError("Code must not be empty.") + + try: + token = load_token(token_file, os.environ) + if session_id is None: + session_id = resolve_session(url=url, token=token, file=file_path) + result = execute_code( + url=url, + session_id=session_id, + token=token, + code=code, + stdout=sys.stdout, + stderr=sys.stderr, + stream=stream, + ) + except PairError as error: + error_text, next_text = _failure_guidance( + error, url=url, session_id=session_id + ) + _emit_failure( + error_text, next_text, session_id=session_id, stream=stream + ) + ctx.exit(2) + except KeyboardInterrupt: + click.echo("Interrupted.", err=True) + ctx.exit(1) + + next_text = None + if not result.success: + next_text = _kernel_next( + result.stderr, url=display_url(url), session_id=session_id + ) -def _codex_skill_dirs() -> list[Path]: - """Return directories where a Codex skill may be installed. + if stream: + if not result.success: + if next_text: + _emit_failure( + "Execution failed.", + next_text, + session_id=session_id, + stream=True, + ) + ctx.exit(1) + return + + payload: dict[str, object] = { + "success": result.success, + "output": result.output, + "stdout": result.stdout, + "stderr": result.stderr, + "session": {"id": session_id}, + } + if next_text: + payload["next"] = next_text + click.echo(json.dumps(payload, indent=2)) + if not result.success: + ctx.exit(1) - Codex loads repository skills from `.agents/skills` directories between - the current directory and repository root. It also loads user and admin - skills from `~/.agents/skills` and `/etc/codex/skills`, respectively. - Keep checking `.codex` for existing direct and plugin installations. - """ - cwd = Path.cwd() - home = Path.home() - roots = [home / ".codex", cwd / ".codex"] - return [ - *_codex_repository_skill_dirs(cwd), - home / ".agents" / "skills", - Path("/etc/codex/skills"), - *[root / "skills" for root in roots], - *[ - skill_dir - for root in roots - for skill_dir in _plugin_skill_dirs(root) - ], - ] +_TOKEN_NEXT = ( + "Pass --token-file or set MARIMO_TOKEN. " + 'Check it with `test -n "$MARIMO_TOKEN"`. Never print the token. ' + "If you have no token, ask the user." +) +_HEADLESS_NEXT = ( + "A headless server has no session until a browser opens the notebook." +) -def _codex_repository_skill_dirs(cwd: Path) -> list[Path]: - """Return Codex skill directories from `cwd` through the repository root.""" - skill_dirs: list[Path] = [] - for directory in (cwd, *cwd.parents): - skill_dirs.append(directory / ".agents" / "skills") - try: - is_repository_root = (directory / ".git").exists() - except OSError: - # Do not search above an ancestor whose repository status cannot - # be determined. The global user and admin paths remain available. - return skill_dirs - if is_repository_root: - return skill_dirs - - # Outside a Git repository, Codex still checks the current directory. - return skill_dirs[:1] - - -def _opencode_skill_dirs() -> list[Path]: - """Return directories where an opencode skill (or compatible layout) may live. - - https://opencode.ai/docs/skills/ - Checked roots are the parent of `/SKILL.md` for: - - - Project opencode: `.opencode/skills/` - - Global opencode: `~/.config/opencode/skills/` - - Project Claude-compatible: `.claude/skills/` - - Global Claude-compatible: `~/.claude/skills/` - - Project agent-compatible: `.agents/skills/` - - Global agent-compatible: `~/.agents/skills/` +def _inspect_command(url: str, session_id: str) -> str: + return ( + f"marimo pair execute --url {url} --session {session_id} " + "--code-file - <<'PY'\n" + "import marimo._code_mode as cm\n" + "async with cm.get_context() as ctx:\n" + " for cell in ctx.cells.values():\n" + " print(cell.id, repr(cell.code), cell.status, cell.errors)\n" + "PY" + ) + + +def _read_cell_command(url: str, session_id: str, cell_id: str) -> str: + return ( + f"marimo pair execute --url {url} --session {session_id} " + "--code-file - <<'PY'\n" + "import marimo._code_mode as cm\n" + "async with cm.get_context() as ctx:\n" + f' cell = ctx.cells["{cell_id}"]\n' + " print(cell.status, cell.errors, " + "[o.data for o in cell.console_outputs])\n" + "PY" + ) + + +def _help_cm_command(url: str, session_id: str) -> str: + return ( + f"marimo pair execute --url {url} --session {session_id} " + "-c 'import marimo._code_mode as cm; help(cm)'" + ) + + +def _kernel_next(stderr: str, *, url: str, session_id: str) -> str | None: + """Recovery guidance for the code-mode mistakes agents make most. + + Matches the traceback text the kernel returned. Unknown failures get + no guidance rather than a guess. """ - cwd = Path.cwd() - home = Path.home() - return [ - cwd / ".opencode" / "skills", - home / ".config" / "opencode" / "skills", - cwd / ".claude" / "skills", - home / ".claude" / "skills", - cwd / ".agents" / "skills", - home / ".agents" / "skills", - ] + if "asyncio.run() cannot be called from a running event loop" in stderr: + return ( + "The kernel already runs an event loop. " + "Write the block at top level:\n" + "import marimo._code_mode as cm\n" + "async with cm.get_context() as ctx:\n" + " ..." + ) + match = re.search(r"KeyError: [\"']Cell '([^']+)' not found", stderr) + if match: + cell_id = match.group(1) + return ( + f"Edits apply when the async with block exits, so '{cell_id}' " + "does not exist yet in this block. " + "Read it in a separate execute:\n" + + _read_cell_command(url, session_id, cell_id) + ) + match = re.search( + r"AttributeError: '?([\w.]+)'? (?:object )?has no attribute '(\w+)'", + stderr, + ) + if match: + owner, name = match.groups() + hint = "" + if owner.endswith("_CellsView"): + hint = " Use ctx.cells[cid] or ctx.cells.values()." + return ( + f"'{owner}' has no attribute '{name}'.{hint} For the full API:\n" + + _help_cm_command(url, session_id) + ) + return None + + +def _failure_guidance( + error: PairError, *, url: str, session_id: str | None +) -> tuple[str, str | None]: + """Error text and one `next` step for a failure before the kernel ran.""" + safe_url = display_url(url) + if isinstance(error, AmbiguousSessionError): + lines = [ + "Pass --session to choose one. Never switch sessions silently." + ] + lines.extend( + f"{session_id}: marimo pair execute --url " + f"{display_url(error.url)} --session {session_id} ..." + for session_id in error.candidates + ) + return str(error), "\n".join(lines) + if isinstance(error, NoSessionError): + return str(error), ( + f"{_HEADLESS_NEXT} Ask the user to open it, then:\n" + f"marimo pair notebook list --url {display_url(error.url)}" + ) + if isinstance(error, StaleSessionError): + return "The session is stale.", ( + "Sessions change when the page reloads. List them again:\n" + f"marimo pair notebook list --url {safe_url}" + ) + message = str(error) + if message == "Authentication failed.": + return message, _TOKEN_NEXT + if message.startswith("The server URL"): + return message, ("Use the url field from:\nmarimo pair notebook list") + if message == "Could not connect to the server.": + return f"Could not connect to {safe_url}.", ( + "Check the URL. To find live servers: marimo pair notebook list" + ) + if message == ( + "The execution response ended before completion was confirmed." + ): + return "Outcome unknown. The code may have run. Do not retry it.", ( + "Inspect the notebook first:\n" + + _inspect_command(safe_url, session_id or "") + ) + return message, None -def pair_agents() -> dict[str, AgentConfig]: - """Return agent configs; paths use `Path.cwd()` at call time.""" - return { - "claude": AgentConfig( - name="Claude Code", - skill_dirs=_claude_skill_dirs(), - ), - "codex": AgentConfig( - name="Codex", - skill_dirs=_codex_skill_dirs(), - ), - "opencode": AgentConfig( - name="opencode", - skill_dirs=_opencode_skill_dirs(), - ), +def _emit_failure( + error_text: str, + next_text: str | None, + *, + session_id: str | None, + stream: bool, +) -> None: + if stream: + click.echo(f"error: {error_text}", err=True) + if next_text: + first, *rest = next_text.splitlines() + click.echo(f" next: {first}", err=True) + for line in rest: + click.echo(f" {line}", err=True) + return + payload: dict[str, object] = { + "success": False, + "error": error_text, } + if next_text: + payload["next"] = next_text + payload.update( + { + "output": None, + "stdout": None, + "stderr": None, + "session": {"id": session_id}, + } + ) + click.echo(json.dumps(payload, indent=2)) -@click.group( - cls=ColoredGroup, - help="""Commands for pair programming with AI.""", +@click.command( + cls=_DocsCommand, + help="Read notebook guidance on demand.", + short_help="""Read notebook guidance on demand.""", ) -def pair() -> None: - pass +@click.argument("topic", required=False) +def docs(topic: str | None) -> None: + topics = _doc_topics() + if topic is None: + for name, title in topics.items(): + click.echo(f"{name} {title}") + return + + if topic not in topics: + valid_topics = ", ".join(topics) + raise click.UsageError( + f"Unknown topic {topic!r}. Valid topics: {valid_topics}." + ) + + click.echo( + (_REFERENCES_DIR / f"{topic}.md").read_text(encoding="utf-8"), + nl=False, + ) @click.command( cls=ColoredCommand, help="""Generate a prompt for pair programming on a running marimo notebook.""", + short_help="""Generate pairing instructions.""", ) @click.option( "--url", @@ -175,6 +429,13 @@ def pair() -> None: type=str, help="URL of the running marimo kernel.", ) +@click.option( + "--session", + "session_id", + required=True, + type=str, + help="Current session ID.", +) @click.option( "--file", "file_path", @@ -182,24 +443,9 @@ def pair() -> None: type=str, help="Notebook path or file key from the page URL.", ) -@click.option( - "--claude", - is_flag=True, - default=False, - help="Validate that the marimo-pair Claude Code skill is installed.", -) -@click.option( - "--codex", - is_flag=True, - default=False, - help="Validate that the marimo-pair Codex skill is installed.", -) -@click.option( - "--opencode", - is_flag=True, - default=False, - help="Validate that the marimo-pair opencode skill is installed.", -) +@click.option("--claude", is_flag=True, hidden=True, expose_value=False) +@click.option("--codex", is_flag=True, hidden=True, expose_value=False) +@click.option("--opencode", is_flag=True, hidden=True, expose_value=False) @click.option( "--with-token", is_flag=True, @@ -208,10 +454,8 @@ def pair() -> None: ) def prompt( url: str, + session_id: str, file_path: str | None, - claude: bool, - codex: bool, - opencode: bool, with_token: bool, ) -> None: """ @@ -219,45 +463,18 @@ def prompt( Example usage: - claude "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --claude)" - codex "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --codex)" - opencode "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --opencode)" + claude "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123')" + codex "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123')" + opencode --prompt "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123')" # Connect to a specific notebook - claude "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --file 'notebooks/example.py' --claude)" + claude "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123' --file 'notebooks/example.py')" # With an auth token - claude "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --claude --with-token)" + claude "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123' --with-token)" """ - # Preserve the file key exactly as supplied. Relative keys are resolved by - # the server workspace and may refer to a remote or non-POSIX filesystem. - # Shell-quote dynamic values because this command is copy-pasted into a - # shell and paths may contain spaces or metacharacters. - file_flag = f" --file {shlex.quote(file_path)}" if file_path else "" - execute_cmd = f"execute-code.sh --url {shlex.quote(url)}{file_flag}" - # Validate that the selected agents have the required skills - selected_agents = { - "claude": claude, - "codex": codex, - "opencode": opencode, - } - for key, agent in pair_agents().items(): - if not selected_agents[key]: - continue - if not agent.has_skill(): - click.echo( - f"The marimo-pair skill for {agent.name} could not be found.\n\n" - "Please install it with:\n\n" - " npx skills add marimo-team/marimo-pair\n\n" - "or\n\n" - " uvx deno -A npm:skills add marimo-team/marimo-pair\n\n" - "More instructions at " - "https://github.com/marimo-team/marimo-pair", - err=True, - ) - # Prompt for token and write it to a temp file if --with-token is set - token_hint = "" + token_file: Path | None = None if with_token: token_dir = _token_dir() url_hash = hashlib.sha256(url.encode()).hexdigest()[:6] @@ -273,24 +490,131 @@ def prompt( finally: os.close(fd) - token_hint = ( - f"\n\nAn auth token is stored at {token_file}. " - f"Pass it via `{execute_cmd} " - f"--token \"$(cat '{token_file}')\"`." - ) - - file_hint = f" (file {file_path})" if file_path else "" + target_lines = [ + "Pair with the live marimo notebook at this target:", + f" Server: {url}", + f" Session: {session_id}", + ] + if file_path: + target_lines.append(f" Notebook: {file_path}") + if token_file: + target_lines.append(f" Token file: {token_file}") # Output the prompt to the wrapper agent CLI click.echo( - "Use the /marimo-pair skill to pair-program on a running " - "marimo notebook.\n\n" - f"Connect to the notebook at: {url}{file_hint}\n\n" - f"Use `{execute_cmd}` from the marimo-pair " - "skill to execute code in the notebook." - f"{token_hint}\n\n" - "Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair." + "\n".join(target_lines) + "\n\n" + "Start with: marimo pair --help\n" + "If marimo is not on your PATH, run it the same way this notebook " + "server was started.\n\n" + "Once connected, run `import marimo as mo; " + 'mo.status.toast("Ready to pair")` to let the user know you are ready.' + ) + + +def _notebook_sort_key(notebook: dict[str, object]) -> tuple[str, str, str]: + server = notebook["server"] + assert isinstance(server, dict) + return ( + str(server.get("url") or ""), + str(notebook.get("name") or ""), + str(notebook.get("path") or ""), ) +def _group_notebooks( + url: str, sessions: dict[str, dict[str, str | None]] +) -> list[dict[str, object]]: + grouped: dict[str | None, dict[str, object]] = {} + for session_id, session in sessions.items(): + key = session.get("path") or session.get("filename") + notebook = grouped.setdefault( + key, + { + "server": {"url": url}, + "name": session.get("filename"), + "path": session.get("path"), + "sessions": [], + }, + ) + notebook_sessions = notebook["sessions"] + assert isinstance(notebook_sessions, list) + notebook_sessions.append({"id": session_id}) + + for notebook in grouped.values(): + notebook_sessions = notebook["sessions"] + assert isinstance(notebook_sessions, list) + notebook_sessions.sort(key=lambda session: session["id"]) + return list(grouped.values()) + + +@click.group( + cls=ColoredGroup, + help="Find active notebooks and their sessions.", + short_help="""Find active notebooks and their sessions.""", +) +def notebook() -> None: + pass + + +@click.command( + name="list", + cls=ColoredCommand, + help="List active notebooks and their session IDs.", + short_help="""List active notebooks and their session IDs.""", +) +@click.option( + "--url", + "urls", + multiple=True, + metavar="URL", + help="Server URL. Repeat to list more than one server.", +) +@click.option( + "--token-file", + type=click.Path(path_type=Path, dir_okay=False), + metavar="PATH", + help="Read the server token from a local file. Otherwise use MARIMO_TOKEN, if set.", +) +def list_notebooks(urls: tuple[str, ...], token_file: Path | None) -> None: + try: + token = load_token(token_file, os.environ) if urls else None + except PairInputError as error: + raise click.UsageError(str(error)) from error + selected_urls = list(urls) if urls else registry_urls() + notebooks: list[dict[str, object]] = [] + warnings: list[str] = [] + + for url in selected_urls: + try: + sessions = list_sessions(url=url, token=token) + except PairError as error: + if isinstance(error, PairInputError) and urls: + raise click.UsageError(str(error)) from error + message = str(error).rstrip(".") + warnings.append( + f"Server {display_url(url)} could not be read: {message}." + ) + else: + notebooks.extend(_group_notebooks(display_url(url), sessions)) + + notebooks.sort(key=_notebook_sort_key) + listing: dict[str, object] = {"notebooks": notebooks, "warnings": warnings} + if any("Authentication failed" in warning for warning in warnings): + listing["next"] = ( + "Pass --token-file or set MARIMO_TOKEN, then run notebook " + "list --url again. Never print the token." + ) + elif not notebooks: + listing["next"] = ( + f"No sessions found. {_HEADLESS_NEXT} If you know the server " + "URL, pass --url. If the server needs a token, pass --token-file " + "or set MARIMO_TOKEN." + ) + click.echo(json.dumps(listing, indent=2)) + + +notebook.add_command(list_notebooks) +pair.add_command(execute) +pair.add_command(docs) +pair.add_command(notebook) pair.add_command(prompt) diff --git a/tests/_cli/test_cli_pair.py b/tests/_cli/test_cli_pair.py index 65ca1b14d9b..138c5592a93 100644 --- a/tests/_cli/test_cli_pair.py +++ b/tests/_cli/test_cli_pair.py @@ -2,20 +2,25 @@ from __future__ import annotations import hashlib +import json import sys from pathlib import Path +from typing import Any from unittest.mock import patch +import pytest from click.testing import CliRunner +from inline_snapshot import snapshot from marimo._cli.cli import main as cli_main -from marimo._cli.pair.commands import ( - AgentConfig, - _codex_repository_skill_dirs, - _codex_skill_dirs, - _opencode_skill_dirs, - _plugin_skill_dirs, - pair_agents, +from marimo._cli.pair import commands +from marimo._cli.pair.client import ( + AmbiguousSessionError, + ExecutionResult, + NoSessionError, + PairError, + PairInputError, + StaleSessionError, ) _runner = CliRunner() @@ -26,36 +31,1255 @@ class TestPairGroup: def test_pair_help(self) -> None: result = _runner.invoke(cli_main, ["pair", "--help"]) + + assert result.exit_code == 0 + assert result.output == snapshot("""\ +Usage: main pair [OPTIONS] COMMAND [ARGS]... + + Pair with a live marimo notebook. + + Workflow: + If you do not have the server URL and session id: + marimo pair notebook list + marimo pair execute --url --session --code-file - <<'PY' + import marimo._code_mode as cm + async with cm.get_context() as ctx: + cid = ctx.create_cell("x = 1") + ctx.run_cell(cid) + PY + marimo pair execute --url --session --code-file - <<'PY' + import marimo._code_mode as cm + async with cm.get_context() as ctx: + cell = ctx.cells[""] + print(cell.status, cell.errors, [o.data for o in cell.console_outputs]) + PY + + Rules: + Cells are the unit of work. The scratchpad is temporary; only cm edits persist. + Cells do not run on creation. Call run_cell after create_cell or edit_cell. + Use async with. Do not await ctx methods. + Session IDs change when the page reloads. If execute reports a stale + session, run notebook list again. + + Code-mode API (this marimo version): + ctx.cells # each has .id .code .status .errors .console_outputs + ctx.create_cell(code) # returns the new cell id + ctx.edit_cell(cid, code) + ctx.run_cell(cid) + ctx.delete_cell(cid) + If a cm call fails, run help(cm): + marimo pair execute --url --session -c 'import marimo._code_mode as cm; help(cm)' + +Options: + -h, --help Show this message and exit. + +Commands: + docs Read notebook guidance on demand. + execute Run Python in a live notebook session. + notebook Find active notebooks and their sessions. + prompt Generate pairing instructions. +""") + + def test_prompt_help(self) -> None: + result = _runner.invoke(cli_main, ["pair", "prompt", "--help"]) + assert result.exit_code == 0 + assert "--url" in result.output + assert "--file" in result.output + assert "--session" in result.output + + +class TestPairExecute: + def test_execute_help_is_offline( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_execute(**kwargs: Any) -> ExecutionResult: + del kwargs + raise AssertionError("execute must not run while showing help") + + monkeypatch.setattr(commands, "execute_code", fail_execute) + result = _runner.invoke( + cli_main, + ["pair", "execute", "--help"], + input="print(1)\n", + ) + + assert result.exit_code == 0 + assert result.output == snapshot("""\ +Usage: main pair execute [OPTIONS] + + Run Python in the selected live notebook kernel's scratchpad. + +Options: + --url URL Server URL. [required] + --session ID Current session ID. Resolved from --file when omitted. + --file PATH Notebook path or file key. Used to resolve --session when + omitted. + --token-file PATH Read the server token from a local file. Otherwise use + MARIMO_TOKEN, if set. + -c TEXT Inline Python. + --code-file PATH Read Python from a UTF-8 file, or from stdin when PATH is + '-'. Supply exactly one input option. + --stream Write stdout and stderr as they arrive. Default: print one + JSON result. + -h, --help Show this message and exit. +""") + + @pytest.mark.parametrize( + "arguments", + [ + [], + ["-c", "print(1)", "--code-file", "code.py"], + ], + ) + def test_execute_requires_exactly_one_input( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + arguments: list[str], + ) -> None: + code_file = tmp_path / "code.py" + code_file.write_text("print(2)", encoding="utf-8") + resolved_arguments = [ + str(code_file) if value == "code.py" else value + for value in arguments + ] + calls: list[dict[str, Any]] = [] + monkeypatch.setattr( + commands, + "execute_code", + lambda **kwargs: calls.append(kwargs), + ) + + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + *resolved_arguments, + ], + input="print('ignored')\n", + ) + + assert result.exit_code == 2 + assert "specify -c or --code-file" in result.output + assert calls == [] + + def test_execute_rejects_empty_code(self) -> None: + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "-c", + "", + ], + ) + + assert result.exit_code == 2 + assert "code must not be empty" in result.output + + def test_execute_reads_code_file( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + code_file = tmp_path / "code.py" + code_file.write_text("print('from file')", encoding="utf-8") + calls: list[dict[str, Any]] = [] + + def fake_execute(**kwargs: Any) -> ExecutionResult: + calls.append(kwargs) + return ExecutionResult( + success=True, output=None, stdout="", stderr="" + ) + + monkeypatch.setattr(commands, "execute_code", fake_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "--code-file", + str(code_file), + ], + ) + + assert result.exit_code == 0 + assert calls[0]["code"] == "print('from file')" + + def test_execute_reads_code_from_stdin( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + calls: list[dict[str, Any]] = [] + + def fake_execute(**kwargs: Any) -> ExecutionResult: + calls.append(kwargs) + return ExecutionResult( + success=True, output=None, stdout="", stderr="" + ) + + monkeypatch.setattr(commands, "execute_code", fake_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "--code-file", + "-", + ], + input='x = 1\nprint("Hi")\n', + ) + + assert result.exit_code == 0 + assert calls[0]["code"] == 'x = 1\nprint("Hi")\n' + + def test_execute_rejects_invalid_utf8_code_file( + self, tmp_path: Path + ) -> None: + code_file = tmp_path / "code.py" + code_file.write_bytes(b"\xff") + + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "--code-file", + str(code_file), + ], + ) + + assert result.exit_code == 2 + assert "error: could not read the code file" in result.output + + def test_execute_rejects_unreadable_code_file( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + code_file = tmp_path / "code.py" + code_file.write_text("print(1)", encoding="utf-8") + + def fail_read_text(self: Path, *, encoding: str) -> str: + del self, encoding + raise OSError("permission denied") + + monkeypatch.setattr(Path, "read_text", fail_read_text) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "--code-file", + str(code_file), + ], + ) + + assert result.exit_code == 2 + assert "error: could not read the code file" in result.output + + def test_execute_success_and_default_streaming( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + calls: list[dict[str, Any]] = [] + + def fake_execute(**kwargs: Any) -> ExecutionResult: + calls.append(kwargs) + return ExecutionResult( + success=True, + output={"mimetype": "text/plain", "data": "done"}, + stdout="", + stderr="", + ) + + monkeypatch.setattr(commands, "execute_code", fake_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 0 + assert calls[0]["url"] == TEST_URL + assert calls[0]["session_id"] == "s_ab12cd" + assert calls[0]["code"] == "print(1)" + assert calls[0]["token"] is None + assert calls[0]["stream"] is False + assert json.loads(result.output) == { + "success": True, + "output": {"mimetype": "text/plain", "data": "done"}, + "stdout": "", + "stderr": "", + "session": {"id": "s_ab12cd"}, + } + + def test_execute_failure_exits_one( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_execution(**kwargs: Any) -> ExecutionResult: + del kwargs + return ExecutionResult( + success=False, + output=None, + stdout="", + stderr="failed", + ) + + monkeypatch.setattr(commands, "execute_code", fail_execution) + + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "-c", + "raise ValueError", + ], + ) + + assert result.exit_code == 1 + assert json.loads(result.output) == { + "success": False, + "output": None, + "stdout": "", + "stderr": "failed", + "session": {"id": "s_ab12cd"}, + } + + @pytest.mark.parametrize( + ("stderr", "expected"), + [ + ( + ( + "RuntimeError: asyncio.run() cannot be called from a " + "running event loop" + ), + "Write the block at top level", + ), + ( + "KeyError: \"Cell 'KBiG' not found. Available cell IDs: [Hbol]\"", + 'cell = ctx.cells["KBiG"]', + ), + ( + "AttributeError: '_CellsView' object has no attribute 'get'", + "help(cm)", + ), + ], + ) + def test_execute_kernel_failure_adds_next( + self, monkeypatch: pytest.MonkeyPatch, stderr: str, expected: str + ) -> None: + def fail_execution(**kwargs: Any) -> ExecutionResult: + del kwargs + return ExecutionResult( + success=False, output=None, stdout="", stderr=stderr + ) + + monkeypatch.setattr(commands, "execute_code", fail_execution) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "-c", + "x", + ], + ) + + assert result.exit_code == 1 + payload = json.loads(result.output) + assert payload["success"] is False + assert payload["stderr"] == stderr + assert expected in payload["next"] + assert ( + "--session s_ab12cd" in payload["next"] + or "top level" in (payload["next"]) + ) + + def test_execute_stream_failure_prints_next_lines( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_execute(**kwargs: Any) -> ExecutionResult: + del kwargs + raise StaleSessionError("Invalid session id: s_old") + + monkeypatch.setattr(commands, "execute_code", fail_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "http://one", + "--session", + "s_old", + "--stream", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 2 + assert result.stdout == "" + assert result.stderr == ( + "error: The session is stale.\n" + " next: Sessions change when the page reloads. List them again:\n" + " marimo pair notebook list --url http://one\n" + ) + + def test_execute_auth_failure_has_token_next( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_execute(**kwargs: Any) -> ExecutionResult: + del kwargs + raise PairError("Authentication failed.") + + monkeypatch.setattr(commands, "execute_code", fail_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "http://one", + "--session", + "s", + "-c", + "1", + ], + ) + + assert result.exit_code == 2 + payload = json.loads(result.output) + assert payload["error"] == "Authentication failed." + assert "MARIMO_TOKEN" in payload["next"] + assert "Never print the token" in payload["next"] + + def test_execute_invalid_url_has_list_next(self) -> None: + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "ftp://x", + "--session", + "s", + "-c", + "1", + ], + ) + + assert result.exit_code == 2 + payload = json.loads(result.output) + assert payload["error"] == "The server URL must use http or https." + assert payload["next"].endswith("marimo pair notebook list") + + def test_execute_reports_pair_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_execute(**kwargs: Any) -> ExecutionResult: + del kwargs + raise PairError("Could not execute code.") + + monkeypatch.setattr(commands, "execute_code", fail_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 2 + assert json.loads(result.output) == { + "success": False, + "error": "Could not execute code.", + "output": None, + "stdout": None, + "stderr": None, + "session": {"id": "s_ab12cd"}, + } + + def test_execute_reports_internal_resolution_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_execute(**kwargs: Any) -> ExecutionResult: + del kwargs + raise PairError("Internal: should not happen after resolution.") + + monkeypatch.setattr(commands, "execute_code", fail_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 2 + payload = json.loads(result.output) + assert payload["success"] is False + assert ( + payload["error"] == "Internal: should not happen after resolution." + ) + assert "next" not in payload + + def test_execute_reports_interrupt( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def interrupt(**kwargs: Any) -> ExecutionResult: + del kwargs + raise KeyboardInterrupt + + monkeypatch.setattr(commands, "execute_code", interrupt) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 1 + assert result.stderr == "Interrupted.\n" + + def test_execute_no_stream_and_token_file( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + token_file = tmp_path / "token.txt" + token_file.write_text("secret", encoding="utf-8") + token_calls: list[Path | None] = [] + execute_calls: list[dict[str, Any]] = [] + + def fake_load_token(path: Path | None, environ: Any) -> str | None: + del environ + token_calls.append(path) + return "secret" + + def fake_execute(**kwargs: Any) -> ExecutionResult: + execute_calls.append(kwargs) + return ExecutionResult( + success=True, output=None, stdout="", stderr="" + ) + + monkeypatch.setattr(commands, "load_token", fake_load_token) + monkeypatch.setattr(commands, "execute_code", fake_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "--token-file", + str(token_file), + "--stream", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 0 + assert token_calls == [token_file] + assert execute_calls[0]["token"] == "secret" + assert execute_calls[0]["stream"] is True + assert result.output == "" + + def test_execute_resolves_one_session_from_file( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + resolve_calls: list[dict[str, Any]] = [] + execute_calls: list[dict[str, Any]] = [] + + def fake_resolve(**kwargs: Any) -> str: + resolve_calls.append(kwargs) + return "s_one" + + def fake_execute(**kwargs: Any) -> ExecutionResult: + execute_calls.append(kwargs) + return ExecutionResult( + success=True, output=None, stdout="", stderr="" + ) + + monkeypatch.setattr(commands, "resolve_session", fake_resolve) + monkeypatch.setattr(commands, "execute_code", fake_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "http://one", + "--file", + "analysis.py", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 0 + assert resolve_calls == [ + {"url": "http://one", "token": None, "file": "analysis.py"} + ] + assert execute_calls[0]["session_id"] == "s_one" + + def test_execute_no_match_exits_two_with_list_next( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fake_resolve(**kwargs: Any) -> str: + del kwargs + raise NoSessionError( + "No running session for notebook 'gone.py' on http://one.", + url="http://user:password@one?access_token=secret", + ) + + monkeypatch.setattr(commands, "resolve_session", fake_resolve) + monkeypatch.setattr( + commands, + "execute_code", + lambda **_kwargs: pytest.fail("execute must not run"), + ) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "http://one", + "--file", + "gone.py", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 2 + payload = json.loads(result.output) + assert payload["success"] is False + assert "No running session for notebook 'gone.py'" in payload["error"] + assert payload["next"].endswith( + "marimo pair notebook list --url http://one" + ) + assert payload["session"] == {"id": None} + assert "password" not in result.output + assert "secret" not in result.output + + def test_execute_two_matches_exits_two_with_candidates( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fake_resolve(**kwargs: Any) -> str: + del kwargs + raise AmbiguousSessionError( + "Notebook 'analysis.py' has 2 running sessions on http://one.", + url="http://user:password@one?access_token=secret", + candidates=("s_a", "s_b"), + ) + + monkeypatch.setattr(commands, "resolve_session", fake_resolve) + monkeypatch.setattr( + commands, + "execute_code", + lambda **_kwargs: pytest.fail("execute must not run"), + ) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "http://one", + "--file", + "analysis.py", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 2 + payload = json.loads(result.output) + assert "has 2 running sessions" in payload["error"] + assert ( + "s_a: marimo pair execute --url http://one --session s_a" + in (payload["next"]) + ) + assert ( + "s_b: marimo pair execute --url http://one --session s_b" + in (payload["next"]) + ) + assert "password" not in result.output + assert "secret" not in result.output + + def test_execute_session_skips_resolve( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + execute_calls: list[dict[str, Any]] = [] + + def fake_execute(**kwargs: Any) -> ExecutionResult: + execute_calls.append(kwargs) + return ExecutionResult( + success=True, output=None, stdout="", stderr="" + ) + + monkeypatch.setattr( + commands, + "resolve_session", + lambda **_kwargs: pytest.fail("resolve must not run"), + ) + monkeypatch.setattr(commands, "execute_code", fake_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "http://user:password@one?access_token=secret", + "--session", + "s_given", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 0 + assert execute_calls[0]["session_id"] == "s_given" + + def test_execute_stale_session_maps_to_targeted_message( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_execute(**kwargs: Any) -> ExecutionResult: + del kwargs + raise StaleSessionError("Invalid session id: s_old") + + monkeypatch.setattr(commands, "execute_code", fail_execute) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "http://user:password@one?access_token=secret", + "--session", + "s_old", + "-c", + "print(1)", + ], + ) + + assert result.exit_code == 2 + payload = json.loads(result.output) + assert payload["error"] == "The session is stale." + assert payload["next"].endswith( + "marimo pair notebook list --url http://one" + ) + assert "password" not in result.output + assert "secret" not in result.output + + +class TestPairDocs: + def test_docs_help_lists_bundled_topics(self) -> None: + result = _runner.invoke(cli_main, ["pair", "docs", "--help"]) + + assert result.exit_code == 0 + assert result.output == snapshot("""\ +Usage: main pair docs [OPTIONS] [TOPIC] + + Read notebook guidance on demand. + +Options: + -h, --help Show this message and exit. + +Available topics: + gotchas Gotchas + notebook-improvements Notebook Improvements + rich-representations Rich Representations +""") + + def test_docs_prints_topic(self) -> None: + reference = commands._REFERENCES_DIR / "gotchas.md" + + result = _runner.invoke(cli_main, ["pair", "docs", "gotchas"]) + + assert result.exit_code == 0 + assert result.output == reference.read_text(encoding="utf-8") + + def test_docs_lists_topics(self) -> None: + result = _runner.invoke(cli_main, ["pair", "docs"]) + + assert result.exit_code == 0 + assert result.output == snapshot("""\ +gotchas Gotchas +notebook-improvements Notebook Improvements +rich-representations Rich Representations +""") + + def test_docs_rejects_unknown_topic(self) -> None: + result = _runner.invoke(cli_main, ["pair", "docs", "nope"]) + + assert result.exit_code == 2 + assert ( + "Valid topics: gotchas, notebook-improvements, " + "rich-representations" in result.output + ) + + def test_docs_help_discovers_new_topic( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + (tmp_path / "extra.md").write_text( + "# Extra Guidance\n\nDetails.\n", encoding="utf-8" + ) + monkeypatch.setattr(commands, "_REFERENCES_DIR", tmp_path) + + result = _runner.invoke(cli_main, ["pair", "docs", "--help"]) + + assert result.exit_code == 0 + assert "extra Extra Guidance" in result.output + + def test_docs_rejects_path_traversal(self) -> None: + result = _runner.invoke(cli_main, ["pair", "docs", "../x"]) + + assert result.exit_code == 2 + assert "Valid topics:" in result.output + + +class TestPairNotebooks: + def test_notebook_help(self) -> None: + result = _runner.invoke(cli_main, ["pair", "notebook", "--help"]) + + assert result.exit_code == 0 + assert result.output == snapshot("""\ +Usage: main pair notebook [OPTIONS] COMMAND [ARGS]... + + Find active notebooks and their sessions. + +Options: + -h, --help Show this message and exit. + +Commands: + list List active notebooks and their session IDs. +""") + + def test_notebook_list_help(self) -> None: + result = _runner.invoke( + cli_main, ["pair", "notebook", "list", "--help"] + ) + + assert result.exit_code == 0 + assert result.output == snapshot("""\ +Usage: main pair notebook list [OPTIONS] + + List active notebooks and their session IDs. + +Options: + --url URL Server URL. Repeat to list more than one server. + --token-file PATH Read the server token from a local file. Otherwise use + MARIMO_TOKEN, if set. + -h, --help Show this message and exit. +""") + + def test_list_groups_sessions_for_one_notebook( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + commands, + "list_sessions", + lambda **_kwargs: { + "session-2": { + "filename": "analysis.py", + "path": "/work/analysis.py", + }, + "session-1": { + "filename": "analysis.py", + "path": "/work/analysis.py", + }, + }, + ) + + result = _runner.invoke( + cli_main, + ["pair", "notebook", "list", "--url", "http://one"], + ) + assert result.exit_code == 0 - assert "pair programming" in result.output.lower() - assert "prompt" in result.output + assert json.loads(result.output) == { + "notebooks": [ + { + "server": {"url": "http://one"}, + "name": "analysis.py", + "path": "/work/analysis.py", + "sessions": [ + {"id": "session-1"}, + {"id": "session-2"}, + ], + } + ], + "warnings": [], + } + + def test_list_rejects_missing_token_file(self, tmp_path: Path) -> None: + result = _runner.invoke( + cli_main, + [ + "pair", + "notebook", + "list", + "--url", + "http://one", + "--token-file", + str(tmp_path / "missing-token"), + ], + ) + + assert result.exit_code == 2 + assert "error: could not read the token file" in result.stderr + assert "Traceback" not in result.stderr + + def test_list_rejects_empty_token_file(self, tmp_path: Path) -> None: + token_file = tmp_path / "token" + token_file.write_text("\n", encoding="utf-8") + + result = _runner.invoke( + cli_main, + [ + "pair", + "notebook", + "list", + "--url", + "http://one", + "--token-file", + str(token_file), + ], + ) + + assert result.exit_code == 2 + assert "error: the token file is empty" in result.stderr + assert "Traceback" not in result.stderr + + def test_list_rejects_invalid_explicit_url(self) -> None: + result = _runner.invoke( + cli_main, + ["pair", "notebook", "list", "--url", "ftp://invalid"], + ) + + assert result.exit_code == 2 + assert "error: the server URL must use http or https" in ( + result.stderr + ) + + def test_list_keeps_same_basename_on_two_urls_separate( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("MARIMO_TOKEN", raising=False) + + def fake_list_sessions( + *, url: str, token: str | None + ) -> dict[str, dict[str, str | None]]: + assert token is None + return { + f"session-{url[-1]}": { + "filename": "analysis.py", + "path": f"/work/{url[-1]}/analysis.py", + } + } + + monkeypatch.setattr(commands, "list_sessions", fake_list_sessions) + result = _runner.invoke( + cli_main, + [ + "pair", + "notebook", + "list", + "--url", + "http://two", + "--url", + "http://one", + ], + ) - def test_prompt_help(self) -> None: - result = _runner.invoke(cli_main, ["pair", "prompt", "--help"]) assert result.exit_code == 0 - assert "--url" in result.output - assert "--claude" in result.output - assert "--codex" in result.output - assert "--opencode" in result.output - assert "--file" in result.output - assert "--session" not in result.output + notebooks = json.loads(result.output)["notebooks"] + assert [notebook["server"]["url"] for notebook in notebooks] == [ + "http://one", + "http://two", + ] + assert len(notebooks) == 2 + + def test_list_sorts_same_basename_by_path( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + commands, + "list_sessions", + lambda **_kwargs: { + "session-b": { + "filename": "analysis.py", + "path": "/work/b/analysis.py", + }, + "session-a": { + "filename": "analysis.py", + "path": "/work/a/analysis.py", + }, + }, + ) + + result = _runner.invoke( + cli_main, + ["pair", "notebook", "list", "--url", "http://one"], + ) + + assert result.exit_code == 0 + notebooks = json.loads(result.output)["notebooks"] + assert [notebook["path"] for notebook in notebooks] == [ + "/work/a/analysis.py", + "/work/b/analysis.py", + ] + + def test_list_preserves_results_when_one_url_fails( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("MARIMO_TOKEN", raising=False) + + def fake_list_sessions( + *, url: str, token: str | None + ) -> dict[str, dict[str, str | None]]: + assert token is None + if url == "http://bad": + raise PairError("Could not connect to http://bad.") + return { + "session-1": { + "filename": "analysis.py", + "path": "/work/analysis.py", + } + } + + monkeypatch.setattr(commands, "list_sessions", fake_list_sessions) + result = _runner.invoke( + cli_main, + [ + "pair", + "notebook", + "list", + "--url", + "http://bad", + "--url", + "http://good", + ], + ) + + assert result.exit_code == 0 + assert json.loads(result.output) == { + "notebooks": [ + { + "server": {"url": "http://good"}, + "name": "analysis.py", + "path": "/work/analysis.py", + "sessions": [{"id": "session-1"}], + } + ], + "warnings": [ + "Server http://bad could not be read: Could not connect to http://bad." + ], + } + + def test_list_redacts_displayed_server_url( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + raw_url = "http://user:password@one/base?access_token=secret#fragment" + calls: list[str] = [] + + def fake_list_sessions( + *, url: str, token: str | None + ) -> dict[str, dict[str, str | None]]: + assert token is None + calls.append(url) + return { + "session-1": { + "filename": "analysis.py", + "path": "/work/analysis.py", + } + } + + monkeypatch.setattr(commands, "list_sessions", fake_list_sessions) + result = _runner.invoke( + cli_main, + ["pair", "notebook", "list", "--url", raw_url], + ) + + assert result.exit_code == 0 + assert calls == [raw_url] + assert json.loads(result.output)["notebooks"][0]["server"] == { + "url": "http://one/base" + } + assert "password" not in result.output + assert "secret" not in result.output + + def test_list_warns_for_invalid_discovered_url( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_list_sessions(**_kwargs: Any) -> None: + raise PairInputError("The server URL must use http or https.") + + monkeypatch.setattr(commands, "registry_urls", lambda: ["ftp://bad"]) + monkeypatch.setattr(commands, "list_sessions", fail_list_sessions) + + result = _runner.invoke(cli_main, ["pair", "notebook", "list"]) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["notebooks"] == [] + assert payload["warnings"] == [ + ( + "Server ftp://bad could not be read: " + "The server URL must use http or https." + ) + ] + assert payload["next"].startswith("No sessions found.") + + def test_list_empty_registry_succeeds( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(commands, "registry_urls", list) + monkeypatch.setattr( + commands, + "list_sessions", + lambda **_kwargs: pytest.fail("No server should be queried"), + ) + + result = _runner.invoke(cli_main, ["pair", "notebook", "list"]) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["notebooks"] == [] + assert payload["warnings"] == [] + assert payload["next"].startswith("No sessions found.") + + def test_list_uses_token_only_for_explicit_urls( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + token_file = tmp_path / "token.txt" + load_calls: list[Path | None] = [] + session_calls: list[tuple[str, str | None]] = [] + + def fake_load_token(path: Path | None, environ: Any) -> str | None: + del environ + load_calls.append(path) + return "secret" + + def fake_list_sessions( + *, url: str, token: str | None + ) -> dict[str, dict[str, str | None]]: + session_calls.append((url, token)) + return {} + + monkeypatch.setattr(commands, "load_token", fake_load_token) + monkeypatch.setattr(commands, "list_sessions", fake_list_sessions) + monkeypatch.setattr( + commands, "registry_urls", lambda: ["http://discovered"] + ) + + discovered = _runner.invoke( + cli_main, + [ + "pair", + "notebook", + "list", + "--token-file", + str(token_file), + ], + ) + explicit = _runner.invoke( + cli_main, + [ + "pair", + "notebook", + "list", + "--url", + "http://explicit", + "--token-file", + str(token_file), + ], + ) + + assert discovered.exit_code == 0 + assert explicit.exit_code == 0 + assert load_calls == [token_file] + assert session_calls == [ + ("http://discovered", None), + ("http://explicit", "secret"), + ] class TestPairPrompt: def test_prompt_requires_url(self) -> None: - result = _runner.invoke(cli_main, ["pair", "prompt"]) + result = _runner.invoke( + cli_main, ["pair", "prompt", "--session", "s_ab12cd"] + ) assert result.exit_code != 0 - def test_prompt_outputs_url(self) -> None: + def test_prompt_requires_session(self) -> None: result = _runner.invoke( cli_main, ["pair", "prompt", "--url", TEST_URL] ) - assert result.exit_code == 0 - assert TEST_URL in result.output - assert "execute-code.sh" in result.output - assert "marimo-pair" in result.output + assert result.exit_code != 0 + assert "--session" in result.output - def test_prompt_with_file(self) -> None: + def test_prompt_outputs_cli_bootstrap(self) -> None: result = _runner.invoke( cli_main, [ @@ -63,121 +1287,62 @@ def test_prompt_with_file(self) -> None: "prompt", "--url", TEST_URL, + "--session", + "s_ab12cd", "--file", "notebooks/example.py", ], ) assert result.exit_code == 0 - assert TEST_URL in result.output - assert "notebooks/example.py" in result.output - assert "--file notebooks/example.py" in result.output + assert result.output == snapshot(f"""\ +Pair with the live marimo notebook at this target: + Server: {TEST_URL} + Session: s_ab12cd + Notebook: notebooks/example.py - def test_prompt_without_file_omits_flag(self) -> None: +Start with: marimo pair --help +If marimo is not on your PATH, run it the same way this notebook server was started. + +Once connected, run `import marimo as mo; mo.status.toast("Ready to pair")` to let the user know you are ready. +""") + + def test_prompt_without_file_omits_notebook(self) -> None: result = _runner.invoke( - cli_main, ["pair", "prompt", "--url", TEST_URL] + cli_main, + [ + "pair", + "prompt", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + ], ) assert result.exit_code == 0 - assert "--file" not in result.output - assert "--session" not in result.output + assert "Notebook:" not in result.output - def test_prompt_rejects_removed_session_option(self) -> None: + @pytest.mark.parametrize( + "agent_flag", ["--claude", "--codex", "--opencode"] + ) + def test_prompt_accepts_legacy_agent_flag_silently( + self, agent_flag: str + ) -> None: result = _runner.invoke( cli_main, - ["pair", "prompt", "--url", TEST_URL, "--session", "s_ab12cd"], + [ + "pair", + "prompt", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + agent_flag, + ], ) - assert result.exit_code != 0 - assert "--session" in result.output - - def test_prompt_shell_quotes_file_paths(self) -> None: - cases = [ - ("relative/path.py", "--file relative/path.py"), - ("/tmp/my notebook.py", "--file '/tmp/my notebook.py'"), - ( - r"C:\Users\Jane Doe\notebook.py", - r"--file 'C:\Users\Jane Doe\notebook.py'", - ), - ( - r"\\server\share\my notebook.py", - r"--file '\\server\share\my notebook.py'", - ), - ( - "notebooks/it's.py", - """--file 'notebooks/it'"'"'s.py'""", - ), - ] - for file_path, expected in cases: - result = _runner.invoke( - cli_main, - [ - "pair", - "prompt", - "--url", - TEST_URL, - "--file", - file_path, - ], - ) - assert result.exit_code == 0 - assert expected in result.output - - def test_prompt_shell_quotes_url_with_metacharacters(self) -> None: - # The execute-code.sh command is meant to be copy-pasted into a shell, - # so a url with metacharacters (`&`) must be quoted so it isn't split. - url = "http://localhost:8000?file=a&b" - result = _runner.invoke(cli_main, ["pair", "prompt", "--url", url]) - assert result.exit_code == 0 - assert f"execute-code.sh --url '{url}'" in result.output - - def test_prompt_skill_missing(self) -> None: - with patch.object(AgentConfig, "has_skill", return_value=False): - for flag in ("--claude", "--codex", "--opencode"): - result = _runner.invoke( - cli_main, - ["pair", "prompt", "--url", TEST_URL, flag], - ) - assert result.exit_code == 0, flag - assert "could not be found" in result.output, flag - - def test_prompt_skill_installed(self) -> None: - with patch.object(AgentConfig, "has_skill", return_value=True): - for flag in ("--claude", "--codex", "--opencode"): - result = _runner.invoke( - cli_main, - ["pair", "prompt", "--url", TEST_URL, flag], - ) - assert result.exit_code == 0, flag - assert TEST_URL in result.output, flag - - def test_prompt_handles_skill_permission_error(self) -> None: - with patch.object(Path, "exists", side_effect=PermissionError): - result = _runner.invoke( - cli_main, - ["pair", "prompt", "--url", TEST_URL, "--codex"], - ) - - assert result.exit_code == 0 - assert "could not be found" in result.output - assert TEST_URL in result.output - - def test_prompt_finds_codex_user_skill(self, tmp_path: Path) -> None: - home = tmp_path / "home" - cwd = tmp_path / "project" - skill = home / ".agents" / "skills" / "marimo-pair" / "SKILL.md" - skill.parent.mkdir(parents=True) - skill.write_text("test") - cwd.mkdir() - - with ( - patch.object(Path, "home", return_value=home), - patch.object(Path, "cwd", return_value=cwd), - ): - result = _runner.invoke( - cli_main, - ["pair", "prompt", "--url", TEST_URL, "--codex"], - ) assert result.exit_code == 0 assert "could not be found" not in result.output + assert "install" not in result.output.lower() class TestPairPromptWithToken: @@ -189,14 +1354,21 @@ def test_with_token_writes_file_and_outputs_prompt( ): result = _runner.invoke( cli_main, - ["pair", "prompt", "--url", TEST_URL, "--with-token"], + [ + "pair", + "prompt", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + "--with-token", + ], input="my-secret-token\n", ) assert result.exit_code == 0 assert TEST_URL in result.output - assert "execute-code.sh" in result.output - assert "token" in result.output.lower() - assert "cat" in result.output + assert "Token file:" in result.output + assert "my-secret-token" not in result.output url_hash = hashlib.sha256(TEST_URL.encode()).hexdigest()[:6] token_file = tmp_path / f"{url_hash}-token.txt" @@ -216,6 +1388,8 @@ def test_with_token_and_file(self, tmp_path: Path) -> None: "prompt", "--url", TEST_URL, + "--session", + "s_ab12cd", "--file", "notebooks/my notebook.py", "--with-token", @@ -223,245 +1397,29 @@ def test_with_token_and_file(self, tmp_path: Path) -> None: input="my-secret-token\n", ) assert result.exit_code == 0 - assert "--file 'notebooks/my notebook.py'" in result.output - # The token hint should target the same file. - assert "--file 'notebooks/my notebook.py' --token" in result.output + assert "Notebook: notebooks/my notebook.py" in result.output + assert "Token file:" in result.output + assert "my-secret-token" not in result.output def test_with_token_still_requires_url(self) -> None: result = _runner.invoke( cli_main, - ["pair", "prompt", "--with-token"], + ["pair", "prompt", "--session", "s_ab12cd", "--with-token"], input="tok\n", ) assert result.exit_code != 0 - def test_with_token_and_agent_flag(self, tmp_path: Path) -> None: - with ( - patch.object(AgentConfig, "has_skill", return_value=True), - patch( - "marimo._cli.pair.commands._token_dir", - return_value=tmp_path, - ), - ): - result = _runner.invoke( - cli_main, - [ - "pair", - "prompt", - "--url", - TEST_URL, - "--claude", - "--with-token", - ], - input="secret\n", - ) - assert result.exit_code == 0 - assert TEST_URL in result.output - assert "token" in result.output.lower() - - def test_with_token_and_skill_missing_fails(self) -> None: - with patch.object(AgentConfig, "has_skill", return_value=False): - result = _runner.invoke( - cli_main, - [ - "pair", - "prompt", - "--url", - TEST_URL, - "--claude", - "--with-token", - ], - input="secret\n", - ) - assert result.exit_code == 0 - assert "could not be found" in result.output - def test_without_token_no_token_hint(self) -> None: result = _runner.invoke( - cli_main, ["pair", "prompt", "--url", TEST_URL] + cli_main, + [ + "pair", + "prompt", + "--url", + TEST_URL, + "--session", + "s_ab12cd", + ], ) assert result.exit_code == 0 - assert "cat" not in result.output - - -class TestOpencodeSkillDirs: - def test_opencode_skill_dirs(self) -> None: - cwd = Path.cwd() - home = Path.home() - assert _opencode_skill_dirs() == [ - cwd / ".opencode" / "skills", - home / ".config" / "opencode" / "skills", - cwd / ".claude" / "skills", - home / ".claude" / "skills", - cwd / ".agents" / "skills", - home / ".agents" / "skills", - ] - - -class TestCodexSkillDirs: - def test_codex_skill_dirs_include_supported_global_locations( - self, tmp_path: Path - ) -> None: - home = tmp_path / "home" - cwd = tmp_path / "project" - cwd.mkdir() - - with ( - patch.object(Path, "home", return_value=home), - patch.object(Path, "cwd", return_value=cwd), - ): - skill_dirs = _codex_skill_dirs() - - assert home / ".agents" / "skills" in skill_dirs - assert Path("/etc/codex/skills") in skill_dirs - - def test_codex_repository_skill_dirs_stop_at_repository_root( - self, tmp_path: Path - ) -> None: - repository = tmp_path / "repository" - cwd = repository / "packages" / "notebooks" - cwd.mkdir(parents=True) - (repository / ".git").mkdir() - - assert _codex_repository_skill_dirs(cwd) == [ - cwd / ".agents" / "skills", - cwd.parent / ".agents" / "skills", - repository / ".agents" / "skills", - ] - - def test_codex_repository_skill_dirs_only_check_cwd_without_repository( - self, tmp_path: Path - ) -> None: - cwd = tmp_path / "notebooks" - cwd.mkdir() - - assert _codex_repository_skill_dirs(cwd) == [ - cwd / ".agents" / "skills" - ] - - def test_codex_repository_skill_dirs_stop_on_permission_error( - self, tmp_path: Path - ) -> None: - cwd = tmp_path / "repository" / "notebooks" - cwd.mkdir(parents=True) - - with patch.object(Path, "exists", side_effect=PermissionError): - assert _codex_repository_skill_dirs(cwd) == [ - cwd / ".agents" / "skills" - ] - - -class TestAgentConfig: - def test_has_skill_true(self, tmp_path: Path) -> None: - skill_dir = tmp_path / "skills" - (skill_dir / "marimo-pair").mkdir(parents=True) - (skill_dir / "marimo-pair" / "SKILL.md").write_text("test") - - agent = AgentConfig(name="test", skill_dirs=[skill_dir]) - assert agent.has_skill() is True - - def test_has_skill_false(self, tmp_path: Path) -> None: - agent = AgentConfig(name="test", skill_dirs=[tmp_path / "nonexistent"]) - assert agent.has_skill() is False - - def test_has_skill_empty_dirs(self) -> None: - agent = AgentConfig(name="test", skill_dirs=[]) - assert agent.has_skill() is False - - def test_has_skill_multiple_dirs_first_match(self, tmp_path: Path) -> None: - dir1 = tmp_path / "a" / "skills" - dir2 = tmp_path / "b" / "skills" - (dir1 / "marimo-pair").mkdir(parents=True) - (dir1 / "marimo-pair" / "SKILL.md").write_text("test") - - agent = AgentConfig(name="test", skill_dirs=[dir1, dir2]) - assert agent.has_skill() is True - - def test_has_skill_multiple_dirs_second_match( - self, tmp_path: Path - ) -> None: - dir1 = tmp_path / "a" / "skills" - dir2 = tmp_path / "b" / "skills" - (dir2 / "marimo-pair").mkdir(parents=True) - (dir2 / "marimo-pair" / "SKILL.md").write_text("test") - - agent = AgentConfig(name="test", skill_dirs=[dir1, dir2]) - assert agent.has_skill() is True - - def test_has_skill_skips_permission_error(self, tmp_path: Path) -> None: - agent = AgentConfig( - name="test", - skill_dirs=[tmp_path / "inaccessible", tmp_path / "installed"], - ) - - with patch.object(Path, "exists", side_effect=[PermissionError, True]): - assert agent.has_skill() is True - - -class TestPluginSkillDirs: - def test_pair_agents_discovers_plugin_skills(self, tmp_path: Path) -> None: - claude_skill_dir = ( - tmp_path - / ".claude" - / "plugins" - / "marketplaces" - / "marimo-pair" - / "skills" - / "marimo-pair" - ) - codex_skill_dir = ( - tmp_path - / ".codex" - / "plugins" - / "cache" - / "marimo-pair" - / "marimo-pair" - / "0.0.18" - / "skills" - / "marimo-pair" - ) - claude_skill_dir.mkdir(parents=True) - codex_skill_dir.mkdir(parents=True) - (claude_skill_dir / "SKILL.md").write_text("test") - (codex_skill_dir / "SKILL.md").write_text("test") - - with ( - patch.object(Path, "home", return_value=tmp_path), - patch.object(Path, "cwd", return_value=tmp_path), - ): - agents = pair_agents() - - assert agents["claude"].has_skill() is True - assert agents["codex"].has_skill() is True - - def test_claude_marketplace_layout(self, tmp_path: Path) -> None: - skill_dir = ( - tmp_path / "plugins" / "marketplaces" / "marimo-pair" / "skills" - ) - (skill_dir / "marimo-pair").mkdir(parents=True) - (skill_dir / "marimo-pair" / "SKILL.md").write_text("test") - - agent = AgentConfig( - name="Claude Code", - skill_dirs=_plugin_skill_dirs(tmp_path), - ) - assert agent.has_skill() is True - - def test_plugin_cache_layout(self, tmp_path: Path) -> None: - skill_dir = ( - tmp_path - / "plugins" - / "cache" - / "marimo-pair" - / "marimo-pair" - / "0.0.18" - / "skills" - ) - (skill_dir / "marimo-pair").mkdir(parents=True) - (skill_dir / "marimo-pair" / "SKILL.md").write_text("test") - - agent = AgentConfig( - name="Codex", - skill_dirs=_plugin_skill_dirs(tmp_path), - ) - assert agent.has_skill() is True + assert "Token file:" not in result.output From 5d8d89e801140ab1fc04d426c22796ab07efd862 Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 16 Sep 2026 10:48:00 -0700 Subject: [PATCH 3/7] feat(frontend): generate a session-aware CLI bootstrap from the pair modal The Pair with an agent modal emits the marimo pair --help bootstrap with the server URL and current session id. The prompt never contains the token. --- .../pair-with-agent-commands.test.ts | 136 ++++++------------ .../actions/pair-with-agent-commands.ts | 50 ++++--- .../editor/actions/pair-with-agent-modal.tsx | 18 +-- 3 files changed, 80 insertions(+), 124 deletions(-) diff --git a/frontend/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts b/frontend/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts index c705e13addb..df2153007da 100644 --- a/frontend/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts +++ b/frontend/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts @@ -1,6 +1,6 @@ /* Copyright 2026 Marimo. All rights reserved. */ -import { afterEach, describe, expect, it, vi } from "vitest"; +import { describe, expect, it } from "vitest"; import { type ConnectionInfo, getFileFromURL, @@ -9,55 +9,19 @@ import { getTerminalCommand, maskToken, } from "../pair-with-agent-commands"; -import { shellQuote } from "@/utils/shell"; +import type { SessionId } from "@/core/kernel/session"; const CONNECTION: ConnectionInfo = { url: "http://localhost:8000", + sessionId: "s_ab12cd" as SessionId, file: "notebooks/example.py", }; const CONNECTION_WITHOUT_FILE: ConnectionInfo = { url: "http://localhost:8000", + sessionId: "s_ab12cd" as SessionId, }; -describe("shellQuote", () => { - it("quotes an empty string", () => { - expect(shellQuote("")).toBe("''"); - }); - - it("leaves shell-safe values untouched", () => { - expect(shellQuote("http://localhost:8000")).toBe("http://localhost:8000"); - expect(shellQuote("notebooks/example.py")).toBe("notebooks/example.py"); - }); - - it("quotes values with shell metacharacters", () => { - expect(shellQuote("http://host:8000?a=1&b=2")).toBe( - "'http://host:8000?a=1&b=2'", - ); - expect(shellQuote("has space")).toBe("'has space'"); - expect(shellQuote("$(rm -rf /)")).toBe("'$(rm -rf /)'"); - }); - - it("escapes embedded single quotes without breaking out", () => { - // Closes the quote, emits a literal ' via "'", then reopens: '"'"' - expect(shellQuote("a'b")).toBe(`'a'"'"'b'`); - }); - - it.each([ - ["/tmp/my notebook.py", "'/tmp/my notebook.py'"], - [ - String.raw`C:\Users\Jane Doe\notebook.py`, - String.raw`'C:\Users\Jane Doe\notebook.py'`, - ], - [ - String.raw`\\server\share\my notebook.py`, - String.raw`'\\server\share\my notebook.py'`, - ], - ])("quotes non-portable path %s as one argument", (path, expected) => { - expect(shellQuote(path)).toBe(expected); - }); -}); - describe("getFileFromURL", () => { it("returns undefined when the file query parameter is absent or empty", () => { expect(getFileFromURL("http://localhost:8000")).toBeUndefined(); @@ -88,31 +52,21 @@ describe("getFileFromURL", () => { }); describe("getMarimoCommand", () => { - afterEach(() => { - vi.unstubAllEnvs(); - }); - - it("uses the local checkout in dev", () => { - vi.stubEnv("DEV", true); + it("uses the current project environment", () => { expect(getMarimoCommand()).toBe("uv run marimo"); }); - - it("uses uvx outside of dev", () => { - vi.stubEnv("DEV", false); - expect(getMarimoCommand()).toBe("uvx marimo@latest"); - }); }); describe("getTerminalCommand", () => { it("includes the url and file for each agent", () => { expect(getTerminalCommand("claude", CONNECTION, false)).toBe( - `claude "$(uv run marimo pair prompt --url http://localhost:8000 --file notebooks/example.py --claude)"`, + `claude "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --file notebooks/example.py)"`, ); expect(getTerminalCommand("codex", CONNECTION, false)).toBe( - `codex "$(uv run marimo pair prompt --url http://localhost:8000 --file notebooks/example.py --codex)"`, + `codex "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --file notebooks/example.py)"`, ); expect(getTerminalCommand("opencode", CONNECTION, false)).toBe( - `opencode --prompt "$(uv run marimo pair prompt --url http://localhost:8000 --file notebooks/example.py --opencode)"`, + `opencode --prompt "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --file notebooks/example.py)"`, ); }); @@ -123,13 +77,17 @@ describe("getTerminalCommand", () => { false, ); expect(command).not.toContain("--file"); - expect(command).not.toContain("--session"); + expect(command).toContain("--session s_ab12cd"); }); it("shell-escapes a url containing metacharacters", () => { const command = getTerminalCommand( "claude", - { url: "http://host:8000?auth=a&b", file: "notebook.py" }, + { + url: "http://host:8000?auth=a&b", + sessionId: CONNECTION.sessionId, + file: "notebook.py", + }, false, ); expect(command).toContain("--url 'http://host:8000?auth=a&b'"); @@ -150,15 +108,15 @@ describe("getTerminalCommand", () => { ])("shell-escapes file path %s", (file, expected) => { const command = getTerminalCommand( "claude", - { url: CONNECTION.url, file }, + { url: CONNECTION.url, sessionId: CONNECTION.sessionId, file }, false, ); expect(command).toContain(expected); }); - it("adds --with-token before the agent flag when requested", () => { + it("adds --with-token when requested", () => { const command = getTerminalCommand("claude", CONNECTION, true); - expect(command).toContain("--with-token --claude"); + expect(command).toContain("--with-token)"); }); it("omits --with-token when not requested", () => { @@ -169,51 +127,49 @@ describe("getTerminalCommand", () => { }); describe("getRawPrompt", () => { - it("references the file-scoped execute-code command", () => { - const prompt = getRawPrompt(CONNECTION, null); + it("does not prefix the bootstrap with a launcher", () => { + const prompt = getRawPrompt(CONNECTION, false); + expect(prompt).toContain("Start with: marimo pair --help"); expect(prompt).toContain( - "execute-code.sh --url http://localhost:8000 --file notebooks/example.py", - ); - expect(prompt).toContain( - "Connect to the notebook at: http://localhost:8000 (file notebooks/example.py)", + "If marimo is not on your PATH, run it the same way this notebook server was started.", ); + expect(prompt).not.toContain("uvx marimo@latest pair --help"); + expect(prompt).not.toContain("uv run marimo pair --help"); }); it("omits file targeting when the page URL has no file", () => { - const prompt = getRawPrompt(CONNECTION_WITHOUT_FILE, null); - expect(prompt).toContain("execute-code.sh --url http://localhost:8000"); - expect(prompt).not.toContain("--file"); - expect(prompt).not.toContain("--session"); + const prompt = getRawPrompt(CONNECTION_WITHOUT_FILE, false); + expect(prompt).toContain(" Session: s_ab12cd"); + expect(prompt).not.toContain(" Notebook:"); }); - it("omits the token hint when there is no token", () => { - const prompt = getRawPrompt(CONNECTION, null); - expect(prompt).not.toContain("--token"); - expect(prompt).not.toContain("auth token"); - }); + it("matches the unauthenticated CLI prompt shape", () => { + expect(getRawPrompt(CONNECTION, false)).toMatchInlineSnapshot(` + "Pair with the live marimo notebook at this target: + Server: http://localhost:8000 + Session: s_ab12cd + Notebook: notebooks/example.py - it("includes a file-scoped token hint when a token is present", () => { - const prompt = getRawPrompt(CONNECTION, "secret-token"); - expect(prompt).toContain( - "execute-code.sh --url http://localhost:8000 --file notebooks/example.py --token secret-token", - ); - }); + Start with: marimo pair --help + If marimo is not on your PATH, run it the same way this notebook server was started. - it("shell-escapes a token containing a single quote", () => { - const prompt = getRawPrompt(CONNECTION, "tok'en"); - expect(prompt).toContain(`--token 'tok'"'"'en'`); + Once connected, run \`import marimo as mo; mo.status.toast("Ready to pair")\` to let the user know you are ready." + `); }); - it("matches the CLI prompt shape", () => { - const prompt = getRawPrompt(CONNECTION, null); - expect(prompt).toMatchInlineSnapshot(` - "Use the /marimo-pair skill to pair-program on a running marimo notebook. + it("directs authenticated users to the token-safe terminal flow", () => { + expect(getRawPrompt(CONNECTION, true)).toMatchInlineSnapshot(` + "Pair with the live marimo notebook at this target: + Server: http://localhost:8000 + Session: s_ab12cd + Notebook: notebooks/example.py - Connect to the notebook at: http://localhost:8000 (file notebooks/example.py) + Start with: marimo pair --help + If marimo is not on your PATH, run it the same way this notebook server was started. - Use \`execute-code.sh --url http://localhost:8000 --file notebooks/example.py\` from the marimo-pair skill to execute code in the notebook. + This notebook uses authentication. Run the terminal command with --with-token and paste its output here instead. - Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair." + Once connected, run \`import marimo as mo; mo.status.toast("Ready to pair")\` to let the user know you are ready." `); }); }); diff --git a/frontend/src/components/editor/actions/pair-with-agent-commands.ts b/frontend/src/components/editor/actions/pair-with-agent-commands.ts index 2949c3fc399..a5e24056001 100644 --- a/frontend/src/components/editor/actions/pair-with-agent-commands.ts +++ b/frontend/src/components/editor/actions/pair-with-agent-commands.ts @@ -2,6 +2,7 @@ import { assertNever } from "@/utils/assertNever"; import { KnownQueryParams } from "@/core/constants"; +import type { SessionId } from "@/core/kernel/session"; import { shellQuote } from "@/utils/shell"; export type AgentTab = "claude" | "codex" | "opencode" | "prompt"; @@ -19,9 +20,9 @@ export const AGENT_LABELS: Record = { export const SKILL_INSTALL = "npx skills add marimo-team/marimo-pair"; -/** How to invoke marimo: from the local checkout in dev, else via uvx. */ +/** Invoke marimo from the same project environment as the notebook server. */ export function getMarimoCommand(): string { - return import.meta.env.DEV ? "uv run marimo" : "uvx marimo@latest"; + return "uv run marimo"; } /** Return the server file key from a page URL, preserving its decoded value. */ @@ -37,6 +38,7 @@ function getFileFlag(file: string | undefined): string { /** Identifies the specific running notebook to pair on. */ export interface ConnectionInfo { url: string; + sessionId: SessionId; /** The server's file key, when the page URL identifies a notebook. */ file?: string; } @@ -47,19 +49,19 @@ export interface ConnectionInfo { */ export function getTerminalCommand( agent: Exclude, - { url, file }: ConnectionInfo, + { url, sessionId, file }: ConnectionInfo, withToken: boolean, ): string { const fileFlag = getFileFlag(file); const tokenFlag = withToken ? " --with-token" : ""; - const base = `${getMarimoCommand()} pair prompt --url ${shellQuote(url)}${fileFlag}${tokenFlag}`; + const base = `${getMarimoCommand()} pair prompt --url ${shellQuote(url)} --session ${shellQuote(sessionId)}${fileFlag}${tokenFlag}`; switch (agent) { case "claude": - return `claude "$(${base} --claude)"`; + return `claude "$(${base})"`; case "codex": - return `codex "$(${base} --codex)"`; + return `codex "$(${base})"`; case "opencode": - return `opencode --prompt "$(${base} --opencode)"`; + return `opencode --prompt "$(${base})"`; default: assertNever(agent); } @@ -71,23 +73,31 @@ export function getTerminalCommand( * an agent behaves the same as the terminal commands. */ export function getRawPrompt( - { url, file }: ConnectionInfo, - token: string | null, + { url, sessionId, file }: ConnectionInfo, + hasToken: boolean, ): string { - const fileFlag = getFileFlag(file); - const fileHint = file ? ` (file ${file})` : ""; - const executeCmd = `execute-code.sh --url ${shellQuote(url)}${fileFlag}`; - const tokenHint = token - ? `\n\nUse this auth token when calling \`execute-code.sh\`: \`${executeCmd} --token ${shellQuote(token)}\`.` - : ""; + const targetLines = [ + "Pair with the live marimo notebook at this target:", + ` Server: ${url}`, + ` Session: ${sessionId}`, + ]; + if (file) { + targetLines.push(` Notebook: ${file}`); + } + return [ - "Use the /marimo-pair skill to pair-program on a running marimo notebook.", - "", - `Connect to the notebook at: ${url}${fileHint}`, + ...targetLines, "", - `Use \`${executeCmd}\` from the marimo-pair skill to execute code in the notebook.${tokenHint}`, + "Start with: marimo pair --help", + "If marimo is not on your PATH, run it the same way this notebook server was started.", + ...(hasToken + ? [ + "", + "This notebook uses authentication. Run the terminal command with --with-token and paste its output here instead.", + ] + : []), "", - "Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair.", + 'Once connected, run `import marimo as mo; mo.status.toast("Ready to pair")` to let the user know you are ready.', ].join("\n"); } diff --git a/frontend/src/components/editor/actions/pair-with-agent-modal.tsx b/frontend/src/components/editor/actions/pair-with-agent-modal.tsx index 602bccc90f9..c2f19eace6c 100644 --- a/frontend/src/components/editor/actions/pair-with-agent-modal.tsx +++ b/frontend/src/components/editor/actions/pair-with-agent-modal.tsx @@ -16,6 +16,7 @@ import { Events } from "@/utils/events"; import { Tooltip } from "@/components/ui/tooltip"; import { asRemoteURL, useRuntimeManager } from "@/core/runtime/config"; import { API } from "@/core/network/api"; +import { getSessionId } from "@/core/kernel/session"; import { AGENT_LABELS, AGENT_TABS, @@ -53,6 +54,7 @@ export const PairWithAgentModal: React.FC<{ const hasToken = Boolean(authToken); const connection: ConnectionInfo = { url: runtimeManager.httpURL.toString(), + sessionId: getSessionId(), file: getFileFromURL(window.location.href), }; @@ -125,21 +127,9 @@ export const PairWithAgentModal: React.FC<{ > - + From e4504fde1cc4f336d9bf1dbd35f1245734680009 Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 16 Sep 2026 10:48:23 -0700 Subject: [PATCH 4/7] test(pair): verify execute against a live server Start a real marimo server in a subprocess and run the CLI against it, so the streaming, cancellation, and failure paths are checked end to end. --- tests/_cli/_pair_server.py | 208 ++++++++++++++++++++++++++++ tests/_cli/test_pair_integration.py | 135 ++++++++++++++++++ 2 files changed, 343 insertions(+) create mode 100644 tests/_cli/_pair_server.py create mode 100644 tests/_cli/test_pair_integration.py diff --git a/tests/_cli/_pair_server.py b/tests/_cli/_pair_server.py new file mode 100644 index 00000000000..b70867910cc --- /dev/null +++ b/tests/_cli/_pair_server.py @@ -0,0 +1,208 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +import json +import socket +import subprocess +import sys +import time +import urllib.request +import uuid +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import websockets.sync.client + +if TYPE_CHECKING: + from collections.abc import Generator + from pathlib import Path + + from websockets.sync.client import ClientConnection + + +@dataclass(frozen=True) +class PairTestServer: + url: str + session_id: str + _process: subprocess.Popen[bytes] + _websocket: ClientConnection + _stderr_path: Path + + def wait_for_kernel( + self, + state: str, + *, + timeout: float = 10.0, + ) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if self._process.poll() is not None: + raise RuntimeError( + "marimo server exited with code " + f"{self._process.returncode}:\n{_tail(self._stderr_path)}" + ) + request = urllib.request.Request( + f"{self.url}/api/kernel/status", + headers={"Marimo-Session-Id": self.session_id}, + ) + try: + with urllib.request.urlopen(request, timeout=1) as response: + current = json.load(response)["state"] + if current == state: + return + except (OSError, KeyError, ValueError): + pass + time.sleep(0.05) + raise TimeoutError( + f"kernel did not become {state!r} within {timeout}s" + ) + + def close(self) -> None: + try: + request = urllib.request.Request( + f"{self.url}/api/kernel/restart_session", + method="POST", + data=b"{}", + headers={ + "Content-Type": "application/json", + "Marimo-Session-Id": self.session_id, + }, + ) + with urllib.request.urlopen(request, timeout=5): + pass + except OSError: + pass + try: + self._websocket.close() + except OSError: + pass + _stop_process(self._process) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + port: int = sock.getsockname()[1] + return port + + +def _tail(path: Path) -> str: + try: + return path.read_text(errors="replace")[-4000:] + except OSError: + return "" + + +def _wait_for_server( + url: str, + process: subprocess.Popen[bytes], + stderr_path: Path, + *, + timeout: float = 15.0, +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"marimo server exited with code {process.returncode}:\n" + f"{_tail(stderr_path)}" + ) + try: + with urllib.request.urlopen(url, timeout=1): + return + except OSError: + time.sleep(0.05) + raise TimeoutError( + f"marimo server at {url} did not start in {timeout}s:\n" + f"{_tail(stderr_path)}" + ) + + +def _wait_for_kernel_ready( + websocket: ClientConnection, + *, + timeout: float = 10.0, +) -> None: + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("kernel did not become ready") + message = json.loads(websocket.recv(timeout=remaining)) + if message.get("op") == "kernel-ready": + return + + +def _stop_process(process: subprocess.Popen[bytes]) -> None: + if process.poll() is not None: + return + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + +def _start_server( + notebook: Path, stderr_path: Path, *, attempts: int = 3 +) -> tuple[subprocess.Popen[bytes], str]: + for attempt in range(attempts): + port = _free_port() + with stderr_path.open("wb") as stderr_file: + process = subprocess.Popen( + [ + sys.executable, + "-m", + "marimo", + "edit", + str(notebook), + "--headless", + "--no-token", + "--no-skew-protection", + "--port", + str(port), + ], + stdout=subprocess.DEVNULL, + stderr=stderr_file, + ) + url = f"http://127.0.0.1:{port}" + try: + _wait_for_server(url, process, stderr_path) + except RuntimeError: + _stop_process(process) + if attempt + 1 == attempts: + raise + continue + return process, url + raise AssertionError("server start attempts must be positive") + + +@contextmanager +def pair_test_server(tmp_path: Path) -> Generator[PairTestServer, None, None]: + notebook = tmp_path / "pair-integration.py" + notebook.write_text("import marimo\napp = marimo.App()\n") + stderr_path = tmp_path / "marimo-stderr.log" + + process, url = _start_server(notebook, stderr_path) + try: + session_id = f"pair_{uuid.uuid4().hex[:8]}" + websocket = websockets.sync.client.connect( + f"{url.replace('http://', 'ws://', 1)}/ws?session_id={session_id}", + open_timeout=5, + ) + _wait_for_kernel_ready(websocket) + server = PairTestServer( + url=url, + session_id=session_id, + _process=process, + _websocket=websocket, + _stderr_path=stderr_path, + ) + try: + yield server + finally: + server.close() + finally: + _stop_process(process) diff --git a/tests/_cli/test_pair_integration.py b/tests/_cli/test_pair_integration.py new file mode 100644 index 00000000000..9a757ca014c --- /dev/null +++ b/tests/_cli/test_pair_integration.py @@ -0,0 +1,135 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +import json +import os +import signal +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING + +import pytest + +from tests._cli._pair_server import PairTestServer, pair_test_server + +if TYPE_CHECKING: + from collections.abc import Generator + + +@pytest.fixture(scope="module") +def server( + tmp_path_factory: pytest.TempPathFactory, +) -> Generator[PairTestServer, None, None]: + with pair_test_server(tmp_path_factory.mktemp("pair")) as running: + yield running + + +def _command(server: PairTestServer, *arguments: str) -> list[str]: + return [ + sys.executable, + "-m", + "marimo", + "pair", + "execute", + "--url", + server.url, + "--session", + server.session_id, + *arguments, + ] + + +def _run( + server: PairTestServer, + *arguments: str, + code_input: str | None = None, + timeout: float = 20, +) -> subprocess.CompletedProcess[str]: + return subprocess.run( + _command(server, *arguments), + input=code_input, + text=True, + capture_output=True, + timeout=timeout, + check=False, + ) + + +def test_streams_output_before_execution_finishes( + server: PairTestServer, +) -> None: + process = subprocess.Popen( + _command( + server, + "--stream", + "-c", + 'import time; print("a"); time.sleep(1); print("b")', + ), + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + try: + assert process.stdout is not None + with ThreadPoolExecutor(max_workers=1) as executor: + first_line = executor.submit(process.stdout.readline).result( + timeout=10 + ) + assert first_line == "a\n" + assert process.poll() is None + stdout, stderr = process.communicate(timeout=10) + finally: + if process.poll() is None: + process.kill() + process.wait() + + assert process.returncode == 0, stderr + assert stdout == "b\n" + + +@pytest.mark.skipif(os.name != "posix", reason="SIGINT requires POSIX") +def test_interrupt_disconnects_and_kernel_recovers( + server: PairTestServer, +) -> None: + process = subprocess.Popen( + _command(server, "-c", "import time; time.sleep(30)"), + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + try: + server.wait_for_kernel("running") + process.send_signal(signal.SIGINT) + stdout, stderr = process.communicate(timeout=10) + finally: + if process.poll() is None: + process.kill() + process.wait() + + assert process.returncode == 1, (stdout, stderr) + assert stderr == "Interrupted.\n" + server.wait_for_kernel("idle") + + recovered = _run(server, "-c", "print('alive')") + assert recovered.returncode == 0, recovered.stderr + payload = json.loads(recovered.stdout) + assert payload["success"] is True + assert payload["stdout"] == "alive\n" + assert payload["output"] is None + assert payload["session"]["id"] == server.session_id + + +def test_missing_input_does_not_read_stdin_or_execute( + server: PairTestServer, +) -> None: + result = _run( + server, + code_input="print('must not run')\n", + timeout=2, + ) + + assert result.returncode == 2 + assert "error: specify -c or --code-file" in result.stderr + assert "must not run" not in result.stdout + server.wait_for_kernel("idle", timeout=1) From 8c48a467c64b463c6c518b29ee2367fc233dfa7c Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 16 Sep 2026 13:14:56 -0700 Subject: [PATCH 5/7] feat(pair): prefer Code Mode for packages and cell hygiene in help The workflow snippet installs through ctx.packages.add, and the rules cover package installs, empty cells, and deletions. Wording follows the marimo-pair skill. --- marimo/_cli/pair/commands.py | 8 +++++++- tests/_cli/test_cli_pair.py | 9 ++++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/marimo/_cli/pair/commands.py b/marimo/_cli/pair/commands.py index 99694607141..c75b01e3167 100644 --- a/marimo/_cli/pair/commands.py +++ b/marimo/_cli/pair/commands.py @@ -74,7 +74,8 @@ def format_epilog( marimo pair execute --url --session --code-file - <<'PY' import marimo._code_mode as cm async with cm.get_context() as ctx: - cid = ctx.create_cell("x = 1") + ctx.packages.add("pandas") + cid = ctx.create_cell("import pandas as pd") ctx.run_cell(cid) PY marimo pair execute --url --session --code-file - <<'PY' @@ -89,6 +90,10 @@ def format_epilog( Cells are the unit of work. The scratchpad is temporary; only cm edits persist. Cells do not run on creation. Call run_cell after create_cell or edit_cell. Use async with. Do not await ctx methods. + Install packages with ctx.packages.add, not uv add or pip. Installs change + the project; confirm when the user did not ask. + If an empty cell exists, edit_cell it instead of creating one. + delete_cell drops the cell's variables. Ask before deleting. Session IDs change when the page reloads. If execute reports a stale session, run notebook list again. @@ -99,6 +104,7 @@ def format_epilog( ctx.edit_cell(cid, code) ctx.run_cell(cid) ctx.delete_cell(cid) + ctx.packages.add("pandas>=2") # queued, installs on exit If a cm call fails, run help(cm): marimo pair execute --url --session -c 'import marimo._code_mode as cm; help(cm)' """, diff --git a/tests/_cli/test_cli_pair.py b/tests/_cli/test_cli_pair.py index 138c5592a93..188d78feab6 100644 --- a/tests/_cli/test_cli_pair.py +++ b/tests/_cli/test_cli_pair.py @@ -33,6 +33,7 @@ def test_pair_help(self) -> None: result = _runner.invoke(cli_main, ["pair", "--help"]) assert result.exit_code == 0 + assert "ctx.packages.add" in result.output assert result.output == snapshot("""\ Usage: main pair [OPTIONS] COMMAND [ARGS]... @@ -44,7 +45,8 @@ def test_pair_help(self) -> None: marimo pair execute --url --session --code-file - <<'PY' import marimo._code_mode as cm async with cm.get_context() as ctx: - cid = ctx.create_cell("x = 1") + ctx.packages.add("pandas") + cid = ctx.create_cell("import pandas as pd") ctx.run_cell(cid) PY marimo pair execute --url --session --code-file - <<'PY' @@ -58,6 +60,10 @@ def test_pair_help(self) -> None: Cells are the unit of work. The scratchpad is temporary; only cm edits persist. Cells do not run on creation. Call run_cell after create_cell or edit_cell. Use async with. Do not await ctx methods. + Install packages with ctx.packages.add, not uv add or pip. Installs change + the project; confirm when the user did not ask. + If an empty cell exists, edit_cell it instead of creating one. + delete_cell drops the cell's variables. Ask before deleting. Session IDs change when the page reloads. If execute reports a stale session, run notebook list again. @@ -67,6 +73,7 @@ def test_pair_help(self) -> None: ctx.edit_cell(cid, code) ctx.run_cell(cid) ctx.delete_cell(cid) + ctx.packages.add("pandas>=2") # queued, installs on exit If a cm call fails, run help(cm): marimo pair execute --url --session -c 'import marimo._code_mode as cm; help(cm)' From dbc2047f8491936ebc4ad474e8094660ac218815 Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 16 Sep 2026 13:15:53 -0700 Subject: [PATCH 6/7] fix(pair): quote recovery commands and treat malformed events as unconfirmed Recovery commands shell-quote the URL and session id and embed the cell id as a Python literal. A malformed execute event now raises the same unconfirmed-completion error as a dropped connection, so the caller inspects instead of retrying. Registry entries on port 443 use https. --- marimo/_cli/pair/client.py | 12 +++++++++++- marimo/_cli/pair/commands.py | 19 ++++++++++++------ tests/_cli/test_cli_pair.py | 35 +++++++++++++++++++++++++++++++++- tests/_cli/test_pair_client.py | 29 +++++++++++++++++++++++++++- 4 files changed, 86 insertions(+), 9 deletions(-) diff --git a/marimo/_cli/pair/client.py b/marimo/_cli/pair/client.py index 5105ff3ddb3..090314b4e9d 100644 --- a/marimo/_cli/pair/client.py +++ b/marimo/_cli/pair/client.py @@ -256,7 +256,15 @@ def write_event(target: TextIO, parts: list[str], value: str) -> None: stdout="".join(stdout_parts), stderr="".join(stderr_parts), ) - except (OSError, http.client.HTTPException) as error: + except ( + OSError, + http.client.HTTPException, + ValueError, + KeyError, + TypeError, + ) as error: + # A malformed event means the outcome is unknown. The code may + # have run, so the caller must inspect rather than retry. raise PairError( "The execution response ended before completion was confirmed." ) from error @@ -388,6 +396,8 @@ def registry_urls() -> list[str]: url_host = format_url_host(host, port, route_bind_all_to_loopback=True) if port == 80: urls.append(f"http://{url_host}{base_url}") + elif port == 443: + urls.append(f"https://{url_host}{base_url}") else: urls.append(f"http://{url_host}:{port}{base_url}") return urls diff --git a/marimo/_cli/pair/commands.py b/marimo/_cli/pair/commands.py index c75b01e3167..6c9f9c99162 100644 --- a/marimo/_cli/pair/commands.py +++ b/marimo/_cli/pair/commands.py @@ -5,6 +5,7 @@ import json import os import re +import shlex import sys from pathlib import Path @@ -252,10 +253,17 @@ def execute( ) +def _execute_prefix(url: str, session_id: str) -> str: + """Shell-safe start of an execute command that targets one session.""" + return ( + "marimo pair execute " + f"--url {shlex.quote(url)} --session {shlex.quote(session_id)}" + ) + + def _inspect_command(url: str, session_id: str) -> str: return ( - f"marimo pair execute --url {url} --session {session_id} " - "--code-file - <<'PY'\n" + f"{_execute_prefix(url, session_id)} --code-file - <<'PY'\n" "import marimo._code_mode as cm\n" "async with cm.get_context() as ctx:\n" " for cell in ctx.cells.values():\n" @@ -266,11 +274,10 @@ def _inspect_command(url: str, session_id: str) -> str: def _read_cell_command(url: str, session_id: str, cell_id: str) -> str: return ( - f"marimo pair execute --url {url} --session {session_id} " - "--code-file - <<'PY'\n" + f"{_execute_prefix(url, session_id)} --code-file - <<'PY'\n" "import marimo._code_mode as cm\n" "async with cm.get_context() as ctx:\n" - f' cell = ctx.cells["{cell_id}"]\n' + f" cell = ctx.cells[{cell_id!r}]\n" " print(cell.status, cell.errors, " "[o.data for o in cell.console_outputs])\n" "PY" @@ -279,7 +286,7 @@ def _read_cell_command(url: str, session_id: str, cell_id: str) -> str: def _help_cm_command(url: str, session_id: str) -> str: return ( - f"marimo pair execute --url {url} --session {session_id} " + f"{_execute_prefix(url, session_id)} " "-c 'import marimo._code_mode as cm; help(cm)'" ) diff --git a/tests/_cli/test_cli_pair.py b/tests/_cli/test_cli_pair.py index 188d78feab6..4b37708786f 100644 --- a/tests/_cli/test_cli_pair.py +++ b/tests/_cli/test_cli_pair.py @@ -397,7 +397,7 @@ def fail_execution(**kwargs: Any) -> ExecutionResult: ), ( "KeyError: \"Cell 'KBiG' not found. Available cell IDs: [Hbol]\"", - 'cell = ctx.cells["KBiG"]', + "cell = ctx.cells['KBiG']", ), ( "AttributeError: '_CellsView' object has no attribute 'get'", @@ -439,6 +439,39 @@ def fail_execution(**kwargs: Any) -> ExecutionResult: or "top level" in (payload["next"]) ) + def test_execute_next_quotes_shell_arguments( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + def fail_execution(**kwargs: Any) -> ExecutionResult: + del kwargs + return ExecutionResult( + success=False, + output=None, + stdout="", + stderr='KeyError: "Cell \'a"b\' not found"', + ) + + monkeypatch.setattr(commands, "execute_code", fail_execution) + result = _runner.invoke( + cli_main, + [ + "pair", + "execute", + "--url", + "http://one/a b", + "--session", + "s'1", + "-c", + "x", + ], + ) + + assert result.exit_code == 1 + payload = json.loads(result.output) + assert "--url 'http://one/a b'" in payload["next"] + assert "--session 's'\"'\"'1'" in payload["next"] + assert "ctx.cells['a\"b']" in payload["next"] + def test_execute_stream_failure_prints_next_lines( self, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/_cli/test_pair_client.py b/tests/_cli/test_pair_client.py index a6f057a3ab3..00d6ce57523 100644 --- a/tests/_cli/test_pair_client.py +++ b/tests/_cli/test_pair_client.py @@ -352,6 +352,33 @@ def test_execute_rejects_missing_done( assert len(calls) == 1 +@pytest.mark.parametrize( + "body", + [ + b"event: stdout\ndata: not json\n\n", + b'event: done\ndata: {"output":null}\n\n', + b"event: stdout\ndata: []\n\n", + ], + ids=["non_json", "done_without_success", "wrong_shape"], +) +def test_execute_treats_malformed_events_as_unconfirmed( + monkeypatch: pytest.MonkeyPatch, body: bytes +) -> None: + response = io.BytesIO(body) + _patch_response(monkeypatch, response) + + with pytest.raises(PairError, match="ended before completion"): + client.execute( + url="http://localhost:2718", + session_id="session-1", + token=None, + code="print(1)", + stdout=io.StringIO(), + stderr=io.StringIO(), + stream=False, + ) + + def test_execute_keeps_buffered_output_after_read_error( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -597,7 +624,7 @@ def test_registry_urls_formats_prefix_and_standard_ports( assert client.registry_urls() == [ "http://localhost/prefix", - "http://localhost:443", + "https://localhost", ] From 8cc74a1b70929ec11adcf1662b6f713a4ce2f2c7 Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 16 Sep 2026 14:43:08 -0700 Subject: [PATCH 7/7] fix(pair): restore skill-based pairing prompt Keep the current Pair with Agent experience stable while the CLI-backed prompt is prepared behind an opt-in flag. Restore agent-specific flags, installed-skill validation, and token handling, and use `uvx marimo@latest` for generated commands so the modal does not assume a local uv project. --- .../pair-with-agent-commands.test.ts | 136 ++++-- .../actions/pair-with-agent-commands.ts | 50 +- .../editor/actions/pair-with-agent-modal.tsx | 18 +- marimo/_cli/pair/commands.py | 235 ++++++++-- tests/_cli/test_cli_pair.py | 429 ++++++++++++++---- 5 files changed, 679 insertions(+), 189 deletions(-) diff --git a/frontend/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts b/frontend/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts index df2153007da..c705e13addb 100644 --- a/frontend/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts +++ b/frontend/src/components/editor/actions/__tests__/pair-with-agent-commands.test.ts @@ -1,6 +1,6 @@ /* Copyright 2026 Marimo. All rights reserved. */ -import { describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { type ConnectionInfo, getFileFromURL, @@ -9,19 +9,55 @@ import { getTerminalCommand, maskToken, } from "../pair-with-agent-commands"; -import type { SessionId } from "@/core/kernel/session"; +import { shellQuote } from "@/utils/shell"; const CONNECTION: ConnectionInfo = { url: "http://localhost:8000", - sessionId: "s_ab12cd" as SessionId, file: "notebooks/example.py", }; const CONNECTION_WITHOUT_FILE: ConnectionInfo = { url: "http://localhost:8000", - sessionId: "s_ab12cd" as SessionId, }; +describe("shellQuote", () => { + it("quotes an empty string", () => { + expect(shellQuote("")).toBe("''"); + }); + + it("leaves shell-safe values untouched", () => { + expect(shellQuote("http://localhost:8000")).toBe("http://localhost:8000"); + expect(shellQuote("notebooks/example.py")).toBe("notebooks/example.py"); + }); + + it("quotes values with shell metacharacters", () => { + expect(shellQuote("http://host:8000?a=1&b=2")).toBe( + "'http://host:8000?a=1&b=2'", + ); + expect(shellQuote("has space")).toBe("'has space'"); + expect(shellQuote("$(rm -rf /)")).toBe("'$(rm -rf /)'"); + }); + + it("escapes embedded single quotes without breaking out", () => { + // Closes the quote, emits a literal ' via "'", then reopens: '"'"' + expect(shellQuote("a'b")).toBe(`'a'"'"'b'`); + }); + + it.each([ + ["/tmp/my notebook.py", "'/tmp/my notebook.py'"], + [ + String.raw`C:\Users\Jane Doe\notebook.py`, + String.raw`'C:\Users\Jane Doe\notebook.py'`, + ], + [ + String.raw`\\server\share\my notebook.py`, + String.raw`'\\server\share\my notebook.py'`, + ], + ])("quotes non-portable path %s as one argument", (path, expected) => { + expect(shellQuote(path)).toBe(expected); + }); +}); + describe("getFileFromURL", () => { it("returns undefined when the file query parameter is absent or empty", () => { expect(getFileFromURL("http://localhost:8000")).toBeUndefined(); @@ -52,21 +88,31 @@ describe("getFileFromURL", () => { }); describe("getMarimoCommand", () => { - it("uses the current project environment", () => { + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it("uses the local checkout in dev", () => { + vi.stubEnv("DEV", true); expect(getMarimoCommand()).toBe("uv run marimo"); }); + + it("uses uvx outside of dev", () => { + vi.stubEnv("DEV", false); + expect(getMarimoCommand()).toBe("uvx marimo@latest"); + }); }); describe("getTerminalCommand", () => { it("includes the url and file for each agent", () => { expect(getTerminalCommand("claude", CONNECTION, false)).toBe( - `claude "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --file notebooks/example.py)"`, + `claude "$(uv run marimo pair prompt --url http://localhost:8000 --file notebooks/example.py --claude)"`, ); expect(getTerminalCommand("codex", CONNECTION, false)).toBe( - `codex "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --file notebooks/example.py)"`, + `codex "$(uv run marimo pair prompt --url http://localhost:8000 --file notebooks/example.py --codex)"`, ); expect(getTerminalCommand("opencode", CONNECTION, false)).toBe( - `opencode --prompt "$(uv run marimo pair prompt --url http://localhost:8000 --session s_ab12cd --file notebooks/example.py)"`, + `opencode --prompt "$(uv run marimo pair prompt --url http://localhost:8000 --file notebooks/example.py --opencode)"`, ); }); @@ -77,17 +123,13 @@ describe("getTerminalCommand", () => { false, ); expect(command).not.toContain("--file"); - expect(command).toContain("--session s_ab12cd"); + expect(command).not.toContain("--session"); }); it("shell-escapes a url containing metacharacters", () => { const command = getTerminalCommand( "claude", - { - url: "http://host:8000?auth=a&b", - sessionId: CONNECTION.sessionId, - file: "notebook.py", - }, + { url: "http://host:8000?auth=a&b", file: "notebook.py" }, false, ); expect(command).toContain("--url 'http://host:8000?auth=a&b'"); @@ -108,15 +150,15 @@ describe("getTerminalCommand", () => { ])("shell-escapes file path %s", (file, expected) => { const command = getTerminalCommand( "claude", - { url: CONNECTION.url, sessionId: CONNECTION.sessionId, file }, + { url: CONNECTION.url, file }, false, ); expect(command).toContain(expected); }); - it("adds --with-token when requested", () => { + it("adds --with-token before the agent flag when requested", () => { const command = getTerminalCommand("claude", CONNECTION, true); - expect(command).toContain("--with-token)"); + expect(command).toContain("--with-token --claude"); }); it("omits --with-token when not requested", () => { @@ -127,49 +169,51 @@ describe("getTerminalCommand", () => { }); describe("getRawPrompt", () => { - it("does not prefix the bootstrap with a launcher", () => { - const prompt = getRawPrompt(CONNECTION, false); - expect(prompt).toContain("Start with: marimo pair --help"); + it("references the file-scoped execute-code command", () => { + const prompt = getRawPrompt(CONNECTION, null); expect(prompt).toContain( - "If marimo is not on your PATH, run it the same way this notebook server was started.", + "execute-code.sh --url http://localhost:8000 --file notebooks/example.py", + ); + expect(prompt).toContain( + "Connect to the notebook at: http://localhost:8000 (file notebooks/example.py)", ); - expect(prompt).not.toContain("uvx marimo@latest pair --help"); - expect(prompt).not.toContain("uv run marimo pair --help"); }); it("omits file targeting when the page URL has no file", () => { - const prompt = getRawPrompt(CONNECTION_WITHOUT_FILE, false); - expect(prompt).toContain(" Session: s_ab12cd"); - expect(prompt).not.toContain(" Notebook:"); + const prompt = getRawPrompt(CONNECTION_WITHOUT_FILE, null); + expect(prompt).toContain("execute-code.sh --url http://localhost:8000"); + expect(prompt).not.toContain("--file"); + expect(prompt).not.toContain("--session"); }); - it("matches the unauthenticated CLI prompt shape", () => { - expect(getRawPrompt(CONNECTION, false)).toMatchInlineSnapshot(` - "Pair with the live marimo notebook at this target: - Server: http://localhost:8000 - Session: s_ab12cd - Notebook: notebooks/example.py + it("omits the token hint when there is no token", () => { + const prompt = getRawPrompt(CONNECTION, null); + expect(prompt).not.toContain("--token"); + expect(prompt).not.toContain("auth token"); + }); - Start with: marimo pair --help - If marimo is not on your PATH, run it the same way this notebook server was started. + it("includes a file-scoped token hint when a token is present", () => { + const prompt = getRawPrompt(CONNECTION, "secret-token"); + expect(prompt).toContain( + "execute-code.sh --url http://localhost:8000 --file notebooks/example.py --token secret-token", + ); + }); - Once connected, run \`import marimo as mo; mo.status.toast("Ready to pair")\` to let the user know you are ready." - `); + it("shell-escapes a token containing a single quote", () => { + const prompt = getRawPrompt(CONNECTION, "tok'en"); + expect(prompt).toContain(`--token 'tok'"'"'en'`); }); - it("directs authenticated users to the token-safe terminal flow", () => { - expect(getRawPrompt(CONNECTION, true)).toMatchInlineSnapshot(` - "Pair with the live marimo notebook at this target: - Server: http://localhost:8000 - Session: s_ab12cd - Notebook: notebooks/example.py + it("matches the CLI prompt shape", () => { + const prompt = getRawPrompt(CONNECTION, null); + expect(prompt).toMatchInlineSnapshot(` + "Use the /marimo-pair skill to pair-program on a running marimo notebook. - Start with: marimo pair --help - If marimo is not on your PATH, run it the same way this notebook server was started. + Connect to the notebook at: http://localhost:8000 (file notebooks/example.py) - This notebook uses authentication. Run the terminal command with --with-token and paste its output here instead. + Use \`execute-code.sh --url http://localhost:8000 --file notebooks/example.py\` from the marimo-pair skill to execute code in the notebook. - Once connected, run \`import marimo as mo; mo.status.toast("Ready to pair")\` to let the user know you are ready." + Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair." `); }); }); diff --git a/frontend/src/components/editor/actions/pair-with-agent-commands.ts b/frontend/src/components/editor/actions/pair-with-agent-commands.ts index a5e24056001..2949c3fc399 100644 --- a/frontend/src/components/editor/actions/pair-with-agent-commands.ts +++ b/frontend/src/components/editor/actions/pair-with-agent-commands.ts @@ -2,7 +2,6 @@ import { assertNever } from "@/utils/assertNever"; import { KnownQueryParams } from "@/core/constants"; -import type { SessionId } from "@/core/kernel/session"; import { shellQuote } from "@/utils/shell"; export type AgentTab = "claude" | "codex" | "opencode" | "prompt"; @@ -20,9 +19,9 @@ export const AGENT_LABELS: Record = { export const SKILL_INSTALL = "npx skills add marimo-team/marimo-pair"; -/** Invoke marimo from the same project environment as the notebook server. */ +/** How to invoke marimo: from the local checkout in dev, else via uvx. */ export function getMarimoCommand(): string { - return "uv run marimo"; + return import.meta.env.DEV ? "uv run marimo" : "uvx marimo@latest"; } /** Return the server file key from a page URL, preserving its decoded value. */ @@ -38,7 +37,6 @@ function getFileFlag(file: string | undefined): string { /** Identifies the specific running notebook to pair on. */ export interface ConnectionInfo { url: string; - sessionId: SessionId; /** The server's file key, when the page URL identifies a notebook. */ file?: string; } @@ -49,19 +47,19 @@ export interface ConnectionInfo { */ export function getTerminalCommand( agent: Exclude, - { url, sessionId, file }: ConnectionInfo, + { url, file }: ConnectionInfo, withToken: boolean, ): string { const fileFlag = getFileFlag(file); const tokenFlag = withToken ? " --with-token" : ""; - const base = `${getMarimoCommand()} pair prompt --url ${shellQuote(url)} --session ${shellQuote(sessionId)}${fileFlag}${tokenFlag}`; + const base = `${getMarimoCommand()} pair prompt --url ${shellQuote(url)}${fileFlag}${tokenFlag}`; switch (agent) { case "claude": - return `claude "$(${base})"`; + return `claude "$(${base} --claude)"`; case "codex": - return `codex "$(${base})"`; + return `codex "$(${base} --codex)"`; case "opencode": - return `opencode --prompt "$(${base})"`; + return `opencode --prompt "$(${base} --opencode)"`; default: assertNever(agent); } @@ -73,31 +71,23 @@ export function getTerminalCommand( * an agent behaves the same as the terminal commands. */ export function getRawPrompt( - { url, sessionId, file }: ConnectionInfo, - hasToken: boolean, + { url, file }: ConnectionInfo, + token: string | null, ): string { - const targetLines = [ - "Pair with the live marimo notebook at this target:", - ` Server: ${url}`, - ` Session: ${sessionId}`, - ]; - if (file) { - targetLines.push(` Notebook: ${file}`); - } - + const fileFlag = getFileFlag(file); + const fileHint = file ? ` (file ${file})` : ""; + const executeCmd = `execute-code.sh --url ${shellQuote(url)}${fileFlag}`; + const tokenHint = token + ? `\n\nUse this auth token when calling \`execute-code.sh\`: \`${executeCmd} --token ${shellQuote(token)}\`.` + : ""; return [ - ...targetLines, + "Use the /marimo-pair skill to pair-program on a running marimo notebook.", + "", + `Connect to the notebook at: ${url}${fileHint}`, "", - "Start with: marimo pair --help", - "If marimo is not on your PATH, run it the same way this notebook server was started.", - ...(hasToken - ? [ - "", - "This notebook uses authentication. Run the terminal command with --with-token and paste its output here instead.", - ] - : []), + `Use \`${executeCmd}\` from the marimo-pair skill to execute code in the notebook.${tokenHint}`, "", - 'Once connected, run `import marimo as mo; mo.status.toast("Ready to pair")` to let the user know you are ready.', + "Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair.", ].join("\n"); } diff --git a/frontend/src/components/editor/actions/pair-with-agent-modal.tsx b/frontend/src/components/editor/actions/pair-with-agent-modal.tsx index c2f19eace6c..602bccc90f9 100644 --- a/frontend/src/components/editor/actions/pair-with-agent-modal.tsx +++ b/frontend/src/components/editor/actions/pair-with-agent-modal.tsx @@ -16,7 +16,6 @@ import { Events } from "@/utils/events"; import { Tooltip } from "@/components/ui/tooltip"; import { asRemoteURL, useRuntimeManager } from "@/core/runtime/config"; import { API } from "@/core/network/api"; -import { getSessionId } from "@/core/kernel/session"; import { AGENT_LABELS, AGENT_TABS, @@ -54,7 +53,6 @@ export const PairWithAgentModal: React.FC<{ const hasToken = Boolean(authToken); const connection: ConnectionInfo = { url: runtimeManager.httpURL.toString(), - sessionId: getSessionId(), file: getFileFromURL(window.location.href), }; @@ -127,9 +125,21 @@ export const PairWithAgentModal: React.FC<{ > - + diff --git a/marimo/_cli/pair/commands.py b/marimo/_cli/pair/commands.py index 6c9f9c99162..d6d77679d32 100644 --- a/marimo/_cli/pair/commands.py +++ b/marimo/_cli/pair/commands.py @@ -7,6 +7,7 @@ import re import shlex import sys +from dataclasses import dataclass, field from pathlib import Path import click @@ -27,6 +28,9 @@ ) from marimo._server.ai.skills import utils as skills_utils +SKILL_NAME = "marimo-pair" +SKILL_FILE = "SKILL.md" + _REFERENCES_DIR = ( Path(skills_utils.__file__).parent / "marimo-pair" / "references" ) @@ -44,6 +48,136 @@ def _token_dir() -> Path: return _cached_token_dir +@dataclass(frozen=True) +class AgentConfig: + name: str + skill_dirs: list[Path] = field(default_factory=list) + + def has_skill(self) -> bool: + for directory in self.skill_dirs: + try: + if (directory / SKILL_NAME / SKILL_FILE).exists(): + return True + except OSError: + # Skill detection is advisory. An inaccessible directory must + # not prevent marimo from generating the pairing prompt. + continue + return False + + +def _claude_skill_dirs() -> list[Path]: + """Return all directories where a Claude Code skill may be installed. + + Skills can be installed directly or bundled in a marketplace plugin in + both the global (`~/.claude`) and local (`.claude`) config directories. + """ + roots = [Path.home() / ".claude", Path.cwd() / ".claude"] + subdirs = ["skills", "plugins", str(Path("plugins") / "marketplaces")] + return [ + *[root / sub for root in roots for sub in subdirs], + *[ + skill_dir + for root in roots + for skill_dir in _plugin_skill_dirs(root) + ], + ] + + +def _plugin_skill_dirs(root: Path) -> list[Path]: + """Return skill directories from marketplace and cached plugins.""" + plugins = root / "plugins" + return [ + *plugins.glob("marketplaces/*/skills"), + *plugins.glob(f"cache/*/{SKILL_NAME}/*/skills"), + ] + + +def _codex_skill_dirs() -> list[Path]: + """Return directories where a Codex skill may be installed. + + Codex loads repository skills from `.agents/skills` directories between + the current directory and repository root. It also loads user and admin + skills from `~/.agents/skills` and `/etc/codex/skills`, respectively. + + Keep checking `.codex` for existing direct and plugin installations. + """ + cwd = Path.cwd() + home = Path.home() + roots = [home / ".codex", cwd / ".codex"] + return [ + *_codex_repository_skill_dirs(cwd), + home / ".agents" / "skills", + Path("/etc/codex/skills"), + *[root / "skills" for root in roots], + *[ + skill_dir + for root in roots + for skill_dir in _plugin_skill_dirs(root) + ], + ] + + +def _codex_repository_skill_dirs(cwd: Path) -> list[Path]: + """Return Codex skill directories from `cwd` through the repository root.""" + skill_dirs: list[Path] = [] + for directory in (cwd, *cwd.parents): + skill_dirs.append(directory / ".agents" / "skills") + try: + is_repository_root = (directory / ".git").exists() + except OSError: + # Do not search above an ancestor whose repository status cannot + # be determined. The global user and admin paths remain available. + return skill_dirs + if is_repository_root: + return skill_dirs + + # Outside a Git repository, Codex still checks the current directory. + return skill_dirs[:1] + + +def _opencode_skill_dirs() -> list[Path]: + """Return directories where an opencode skill (or compatible layout) may live. + + https://opencode.ai/docs/skills/ + Checked roots are the parent of `/SKILL.md` for: + + - Project opencode: `.opencode/skills/` + - Global opencode: `~/.config/opencode/skills/` + - Project Claude-compatible: `.claude/skills/` + - Global Claude-compatible: `~/.claude/skills/` + - Project agent-compatible: `.agents/skills/` + - Global agent-compatible: `~/.agents/skills/` + """ + cwd = Path.cwd() + home = Path.home() + return [ + cwd / ".opencode" / "skills", + home / ".config" / "opencode" / "skills", + cwd / ".claude" / "skills", + home / ".claude" / "skills", + cwd / ".agents" / "skills", + home / ".agents" / "skills", + ] + + +def pair_agents() -> dict[str, AgentConfig]: + """Return agent configs; paths use `Path.cwd()` at call time.""" + return { + "claude": AgentConfig( + name="Claude Code", + skill_dirs=_claude_skill_dirs(), + ), + "codex": AgentConfig( + name="Codex", + skill_dirs=_codex_skill_dirs(), + ), + "opencode": AgentConfig( + name="opencode", + skill_dirs=_opencode_skill_dirs(), + ), + } + + def _doc_topics() -> dict[str, str]: return { path.stem: path.read_text(encoding="utf-8") @@ -434,7 +568,6 @@ def docs(topic: str | None) -> None: @click.command( cls=ColoredCommand, help="""Generate a prompt for pair programming on a running marimo notebook.""", - short_help="""Generate pairing instructions.""", ) @click.option( "--url", @@ -442,13 +575,6 @@ def docs(topic: str | None) -> None: type=str, help="URL of the running marimo kernel.", ) -@click.option( - "--session", - "session_id", - required=True, - type=str, - help="Current session ID.", -) @click.option( "--file", "file_path", @@ -456,9 +582,24 @@ def docs(topic: str | None) -> None: type=str, help="Notebook path or file key from the page URL.", ) -@click.option("--claude", is_flag=True, hidden=True, expose_value=False) -@click.option("--codex", is_flag=True, hidden=True, expose_value=False) -@click.option("--opencode", is_flag=True, hidden=True, expose_value=False) +@click.option( + "--claude", + is_flag=True, + default=False, + help="Validate that the marimo-pair Claude Code skill is installed.", +) +@click.option( + "--codex", + is_flag=True, + default=False, + help="Validate that the marimo-pair Codex skill is installed.", +) +@click.option( + "--opencode", + is_flag=True, + default=False, + help="Validate that the marimo-pair opencode skill is installed.", +) @click.option( "--with-token", is_flag=True, @@ -467,8 +608,10 @@ def docs(topic: str | None) -> None: ) def prompt( url: str, - session_id: str, file_path: str | None, + claude: bool, + codex: bool, + opencode: bool, with_token: bool, ) -> None: """ @@ -476,18 +619,45 @@ def prompt( Example usage: - claude "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123')" - codex "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123')" - opencode --prompt "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123')" + claude "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --claude)" + codex "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --codex)" + opencode "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --opencode)" # Connect to a specific notebook - claude "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123' --file 'notebooks/example.py')" + claude "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --file 'notebooks/example.py' --claude)" # With an auth token - claude "$(marimo pair prompt --url 'https://localhost:8000' --session 'session-123' --with-token)" + claude "$(uvx marimo@latest pair prompt --url 'https://localhost:8000' --claude --with-token)" """ + # Preserve the file key exactly as supplied. Relative keys are resolved by + # the server workspace and may refer to a remote or non-POSIX filesystem. + # Shell-quote dynamic values because this command is copy-pasted into a + # shell and paths may contain spaces or metacharacters. + file_flag = f" --file {shlex.quote(file_path)}" if file_path else "" + execute_cmd = f"execute-code.sh --url {shlex.quote(url)}{file_flag}" + # Validate that the selected agents have the required skills + selected_agents = { + "claude": claude, + "codex": codex, + "opencode": opencode, + } + for key, agent in pair_agents().items(): + if not selected_agents[key]: + continue + if not agent.has_skill(): + click.echo( + f"The marimo-pair skill for {agent.name} could not be found.\n\n" + "Please install it with:\n\n" + " npx skills add marimo-team/marimo-pair\n\n" + "or\n\n" + " uvx deno -A npm:skills add marimo-team/marimo-pair\n\n" + "More instructions at " + "https://github.com/marimo-team/marimo-pair", + err=True, + ) + # Prompt for token and write it to a temp file if --with-token is set - token_file: Path | None = None + token_hint = "" if with_token: token_dir = _token_dir() url_hash = hashlib.sha256(url.encode()).hexdigest()[:6] @@ -503,24 +673,23 @@ def prompt( finally: os.close(fd) - target_lines = [ - "Pair with the live marimo notebook at this target:", - f" Server: {url}", - f" Session: {session_id}", - ] - if file_path: - target_lines.append(f" Notebook: {file_path}") - if token_file: - target_lines.append(f" Token file: {token_file}") + token_hint = ( + f"\n\nAn auth token is stored at {token_file}. " + f"Pass it via `{execute_cmd} " + f"--token \"$(cat '{token_file}')\"`." + ) + + file_hint = f" (file {file_path})" if file_path else "" # Output the prompt to the wrapper agent CLI click.echo( - "\n".join(target_lines) + "\n\n" - "Start with: marimo pair --help\n" - "If marimo is not on your PATH, run it the same way this notebook " - "server was started.\n\n" - "Once connected, run `import marimo as mo; " - 'mo.status.toast("Ready to pair")` to let the user know you are ready.' + "Use the /marimo-pair skill to pair-program on a running " + "marimo notebook.\n\n" + f"Connect to the notebook at: {url}{file_hint}\n\n" + f"Use `{execute_cmd}` from the marimo-pair " + "skill to execute code in the notebook." + f"{token_hint}\n\n" + "Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo letting them know you're ready to pair." ) diff --git a/tests/_cli/test_cli_pair.py b/tests/_cli/test_cli_pair.py index 4b37708786f..2ac4735aec7 100644 --- a/tests/_cli/test_cli_pair.py +++ b/tests/_cli/test_cli_pair.py @@ -22,6 +22,14 @@ PairInputError, StaleSessionError, ) +from marimo._cli.pair.commands import ( + AgentConfig, + _codex_repository_skill_dirs, + _codex_skill_dirs, + _opencode_skill_dirs, + _plugin_skill_dirs, + pair_agents, +) _runner = CliRunner() @@ -84,15 +92,18 @@ def test_pair_help(self) -> None: docs Read notebook guidance on demand. execute Run Python in a live notebook session. notebook Find active notebooks and their sessions. - prompt Generate pairing instructions. + prompt Generate a prompt for pair programming on... """) def test_prompt_help(self) -> None: result = _runner.invoke(cli_main, ["pair", "prompt", "--help"]) assert result.exit_code == 0 assert "--url" in result.output + assert "--claude" in result.output + assert "--codex" in result.output + assert "--opencode" in result.output assert "--file" in result.output - assert "--session" in result.output + assert "--session" not in result.output class TestPairExecute: @@ -1307,19 +1318,19 @@ def fake_list_sessions( class TestPairPrompt: def test_prompt_requires_url(self) -> None: - result = _runner.invoke( - cli_main, ["pair", "prompt", "--session", "s_ab12cd"] - ) + result = _runner.invoke(cli_main, ["pair", "prompt"]) assert result.exit_code != 0 - def test_prompt_requires_session(self) -> None: + def test_prompt_outputs_url(self) -> None: result = _runner.invoke( cli_main, ["pair", "prompt", "--url", TEST_URL] ) - assert result.exit_code != 0 - assert "--session" in result.output + assert result.exit_code == 0 + assert TEST_URL in result.output + assert "execute-code.sh" in result.output + assert "marimo-pair" in result.output - def test_prompt_outputs_cli_bootstrap(self) -> None: + def test_prompt_with_file(self) -> None: result = _runner.invoke( cli_main, [ @@ -1327,62 +1338,121 @@ def test_prompt_outputs_cli_bootstrap(self) -> None: "prompt", "--url", TEST_URL, - "--session", - "s_ab12cd", "--file", "notebooks/example.py", ], ) assert result.exit_code == 0 - assert result.output == snapshot(f"""\ -Pair with the live marimo notebook at this target: - Server: {TEST_URL} - Session: s_ab12cd - Notebook: notebooks/example.py - -Start with: marimo pair --help -If marimo is not on your PATH, run it the same way this notebook server was started. - -Once connected, run `import marimo as mo; mo.status.toast("Ready to pair")` to let the user know you are ready. -""") + assert TEST_URL in result.output + assert "notebooks/example.py" in result.output + assert "--file notebooks/example.py" in result.output - def test_prompt_without_file_omits_notebook(self) -> None: + def test_prompt_without_file_omits_flag(self) -> None: result = _runner.invoke( - cli_main, - [ - "pair", - "prompt", - "--url", - TEST_URL, - "--session", - "s_ab12cd", - ], + cli_main, ["pair", "prompt", "--url", TEST_URL] ) assert result.exit_code == 0 - assert "Notebook:" not in result.output + assert "--file" not in result.output + assert "--session" not in result.output - @pytest.mark.parametrize( - "agent_flag", ["--claude", "--codex", "--opencode"] - ) - def test_prompt_accepts_legacy_agent_flag_silently( - self, agent_flag: str - ) -> None: + def test_prompt_rejects_removed_session_option(self) -> None: result = _runner.invoke( cli_main, - [ - "pair", - "prompt", - "--url", - TEST_URL, - "--session", - "s_ab12cd", - agent_flag, - ], + ["pair", "prompt", "--url", TEST_URL, "--session", "s_ab12cd"], ) + assert result.exit_code != 0 + assert "--session" in result.output + + def test_prompt_shell_quotes_file_paths(self) -> None: + cases = [ + ("relative/path.py", "--file relative/path.py"), + ("/tmp/my notebook.py", "--file '/tmp/my notebook.py'"), + ( + r"C:\Users\Jane Doe\notebook.py", + r"--file 'C:\Users\Jane Doe\notebook.py'", + ), + ( + r"\\server\share\my notebook.py", + r"--file '\\server\share\my notebook.py'", + ), + ( + "notebooks/it's.py", + """--file 'notebooks/it'"'"'s.py'""", + ), + ] + for file_path, expected in cases: + result = _runner.invoke( + cli_main, + [ + "pair", + "prompt", + "--url", + TEST_URL, + "--file", + file_path, + ], + ) + assert result.exit_code == 0 + assert expected in result.output + + def test_prompt_shell_quotes_url_with_metacharacters(self) -> None: + # The execute-code.sh command is meant to be copy-pasted into a shell, + # so a url with metacharacters (`&`) must be quoted so it isn't split. + url = "http://localhost:8000?file=a&b" + result = _runner.invoke(cli_main, ["pair", "prompt", "--url", url]) + assert result.exit_code == 0 + assert f"execute-code.sh --url '{url}'" in result.output + + def test_prompt_skill_missing(self) -> None: + with patch.object(AgentConfig, "has_skill", return_value=False): + for flag in ("--claude", "--codex", "--opencode"): + result = _runner.invoke( + cli_main, + ["pair", "prompt", "--url", TEST_URL, flag], + ) + assert result.exit_code == 0, flag + assert "could not be found" in result.output, flag + + def test_prompt_skill_installed(self) -> None: + with patch.object(AgentConfig, "has_skill", return_value=True): + for flag in ("--claude", "--codex", "--opencode"): + result = _runner.invoke( + cli_main, + ["pair", "prompt", "--url", TEST_URL, flag], + ) + assert result.exit_code == 0, flag + assert TEST_URL in result.output, flag + + def test_prompt_handles_skill_permission_error(self) -> None: + with patch.object(Path, "exists", side_effect=PermissionError): + result = _runner.invoke( + cli_main, + ["pair", "prompt", "--url", TEST_URL, "--codex"], + ) + + assert result.exit_code == 0 + assert "could not be found" in result.output + assert TEST_URL in result.output + + def test_prompt_finds_codex_user_skill(self, tmp_path: Path) -> None: + home = tmp_path / "home" + cwd = tmp_path / "project" + skill = home / ".agents" / "skills" / "marimo-pair" / "SKILL.md" + skill.parent.mkdir(parents=True) + skill.write_text("test") + cwd.mkdir() + + with ( + patch.object(Path, "home", return_value=home), + patch.object(Path, "cwd", return_value=cwd), + ): + result = _runner.invoke( + cli_main, + ["pair", "prompt", "--url", TEST_URL, "--codex"], + ) assert result.exit_code == 0 assert "could not be found" not in result.output - assert "install" not in result.output.lower() class TestPairPromptWithToken: @@ -1394,21 +1464,14 @@ def test_with_token_writes_file_and_outputs_prompt( ): result = _runner.invoke( cli_main, - [ - "pair", - "prompt", - "--url", - TEST_URL, - "--session", - "s_ab12cd", - "--with-token", - ], + ["pair", "prompt", "--url", TEST_URL, "--with-token"], input="my-secret-token\n", ) assert result.exit_code == 0 assert TEST_URL in result.output - assert "Token file:" in result.output - assert "my-secret-token" not in result.output + assert "execute-code.sh" in result.output + assert "token" in result.output.lower() + assert "cat" in result.output url_hash = hashlib.sha256(TEST_URL.encode()).hexdigest()[:6] token_file = tmp_path / f"{url_hash}-token.txt" @@ -1428,8 +1491,6 @@ def test_with_token_and_file(self, tmp_path: Path) -> None: "prompt", "--url", TEST_URL, - "--session", - "s_ab12cd", "--file", "notebooks/my notebook.py", "--with-token", @@ -1437,29 +1498,245 @@ def test_with_token_and_file(self, tmp_path: Path) -> None: input="my-secret-token\n", ) assert result.exit_code == 0 - assert "Notebook: notebooks/my notebook.py" in result.output - assert "Token file:" in result.output - assert "my-secret-token" not in result.output + assert "--file 'notebooks/my notebook.py'" in result.output + # The token hint should target the same file. + assert "--file 'notebooks/my notebook.py' --token" in result.output def test_with_token_still_requires_url(self) -> None: result = _runner.invoke( cli_main, - ["pair", "prompt", "--session", "s_ab12cd", "--with-token"], + ["pair", "prompt", "--with-token"], input="tok\n", ) assert result.exit_code != 0 + def test_with_token_and_agent_flag(self, tmp_path: Path) -> None: + with ( + patch.object(AgentConfig, "has_skill", return_value=True), + patch( + "marimo._cli.pair.commands._token_dir", + return_value=tmp_path, + ), + ): + result = _runner.invoke( + cli_main, + [ + "pair", + "prompt", + "--url", + TEST_URL, + "--claude", + "--with-token", + ], + input="secret\n", + ) + assert result.exit_code == 0 + assert TEST_URL in result.output + assert "token" in result.output.lower() + + def test_with_token_and_skill_missing_fails(self) -> None: + with patch.object(AgentConfig, "has_skill", return_value=False): + result = _runner.invoke( + cli_main, + [ + "pair", + "prompt", + "--url", + TEST_URL, + "--claude", + "--with-token", + ], + input="secret\n", + ) + assert result.exit_code == 0 + assert "could not be found" in result.output + def test_without_token_no_token_hint(self) -> None: result = _runner.invoke( - cli_main, - [ - "pair", - "prompt", - "--url", - TEST_URL, - "--session", - "s_ab12cd", - ], + cli_main, ["pair", "prompt", "--url", TEST_URL] ) assert result.exit_code == 0 - assert "Token file:" not in result.output + assert "cat" not in result.output + + +class TestOpencodeSkillDirs: + def test_opencode_skill_dirs(self) -> None: + cwd = Path.cwd() + home = Path.home() + assert _opencode_skill_dirs() == [ + cwd / ".opencode" / "skills", + home / ".config" / "opencode" / "skills", + cwd / ".claude" / "skills", + home / ".claude" / "skills", + cwd / ".agents" / "skills", + home / ".agents" / "skills", + ] + + +class TestCodexSkillDirs: + def test_codex_skill_dirs_include_supported_global_locations( + self, tmp_path: Path + ) -> None: + home = tmp_path / "home" + cwd = tmp_path / "project" + cwd.mkdir() + + with ( + patch.object(Path, "home", return_value=home), + patch.object(Path, "cwd", return_value=cwd), + ): + skill_dirs = _codex_skill_dirs() + + assert home / ".agents" / "skills" in skill_dirs + assert Path("/etc/codex/skills") in skill_dirs + + def test_codex_repository_skill_dirs_stop_at_repository_root( + self, tmp_path: Path + ) -> None: + repository = tmp_path / "repository" + cwd = repository / "packages" / "notebooks" + cwd.mkdir(parents=True) + (repository / ".git").mkdir() + + assert _codex_repository_skill_dirs(cwd) == [ + cwd / ".agents" / "skills", + cwd.parent / ".agents" / "skills", + repository / ".agents" / "skills", + ] + + def test_codex_repository_skill_dirs_only_check_cwd_without_repository( + self, tmp_path: Path + ) -> None: + cwd = tmp_path / "notebooks" + cwd.mkdir() + + assert _codex_repository_skill_dirs(cwd) == [ + cwd / ".agents" / "skills" + ] + + def test_codex_repository_skill_dirs_stop_on_permission_error( + self, tmp_path: Path + ) -> None: + cwd = tmp_path / "repository" / "notebooks" + cwd.mkdir(parents=True) + + with patch.object(Path, "exists", side_effect=PermissionError): + assert _codex_repository_skill_dirs(cwd) == [ + cwd / ".agents" / "skills" + ] + + +class TestAgentConfig: + def test_has_skill_true(self, tmp_path: Path) -> None: + skill_dir = tmp_path / "skills" + (skill_dir / "marimo-pair").mkdir(parents=True) + (skill_dir / "marimo-pair" / "SKILL.md").write_text("test") + + agent = AgentConfig(name="test", skill_dirs=[skill_dir]) + assert agent.has_skill() is True + + def test_has_skill_false(self, tmp_path: Path) -> None: + agent = AgentConfig(name="test", skill_dirs=[tmp_path / "nonexistent"]) + assert agent.has_skill() is False + + def test_has_skill_empty_dirs(self) -> None: + agent = AgentConfig(name="test", skill_dirs=[]) + assert agent.has_skill() is False + + def test_has_skill_multiple_dirs_first_match(self, tmp_path: Path) -> None: + dir1 = tmp_path / "a" / "skills" + dir2 = tmp_path / "b" / "skills" + (dir1 / "marimo-pair").mkdir(parents=True) + (dir1 / "marimo-pair" / "SKILL.md").write_text("test") + + agent = AgentConfig(name="test", skill_dirs=[dir1, dir2]) + assert agent.has_skill() is True + + def test_has_skill_multiple_dirs_second_match( + self, tmp_path: Path + ) -> None: + dir1 = tmp_path / "a" / "skills" + dir2 = tmp_path / "b" / "skills" + (dir2 / "marimo-pair").mkdir(parents=True) + (dir2 / "marimo-pair" / "SKILL.md").write_text("test") + + agent = AgentConfig(name="test", skill_dirs=[dir1, dir2]) + assert agent.has_skill() is True + + def test_has_skill_skips_permission_error(self, tmp_path: Path) -> None: + agent = AgentConfig( + name="test", + skill_dirs=[tmp_path / "inaccessible", tmp_path / "installed"], + ) + + with patch.object(Path, "exists", side_effect=[PermissionError, True]): + assert agent.has_skill() is True + + +class TestPluginSkillDirs: + def test_pair_agents_discovers_plugin_skills(self, tmp_path: Path) -> None: + claude_skill_dir = ( + tmp_path + / ".claude" + / "plugins" + / "marketplaces" + / "marimo-pair" + / "skills" + / "marimo-pair" + ) + codex_skill_dir = ( + tmp_path + / ".codex" + / "plugins" + / "cache" + / "marimo-pair" + / "marimo-pair" + / "0.0.18" + / "skills" + / "marimo-pair" + ) + claude_skill_dir.mkdir(parents=True) + codex_skill_dir.mkdir(parents=True) + (claude_skill_dir / "SKILL.md").write_text("test") + (codex_skill_dir / "SKILL.md").write_text("test") + + with ( + patch.object(Path, "home", return_value=tmp_path), + patch.object(Path, "cwd", return_value=tmp_path), + ): + agents = pair_agents() + + assert agents["claude"].has_skill() is True + assert agents["codex"].has_skill() is True + + def test_claude_marketplace_layout(self, tmp_path: Path) -> None: + skill_dir = ( + tmp_path / "plugins" / "marketplaces" / "marimo-pair" / "skills" + ) + (skill_dir / "marimo-pair").mkdir(parents=True) + (skill_dir / "marimo-pair" / "SKILL.md").write_text("test") + + agent = AgentConfig( + name="Claude Code", + skill_dirs=_plugin_skill_dirs(tmp_path), + ) + assert agent.has_skill() is True + + def test_plugin_cache_layout(self, tmp_path: Path) -> None: + skill_dir = ( + tmp_path + / "plugins" + / "cache" + / "marimo-pair" + / "marimo-pair" + / "0.0.18" + / "skills" + ) + (skill_dir / "marimo-pair").mkdir(parents=True) + (skill_dir / "marimo-pair" / "SKILL.md").write_text("test") + + agent = AgentConfig( + name="Codex", + skill_dirs=_plugin_skill_dirs(tmp_path), + ) + assert agent.has_skill() is True