diff --git a/.woodpecker/check.yml b/.woodpecker/check.yml deleted file mode 100644 index 6f7273b..0000000 --- a/.woodpecker/check.yml +++ /dev/null @@ -1,28 +0,0 @@ -steps: - unit-tests: - image: ghcr.io/astral-sh/uv:python3.13-bookworm - when: - event: push - path: - - "src/**" - - "tests/**" - commands: - - uv sync --dev - - uv run pytest -m "not integration" -v - - integration-tests: - image: ghcr.io/astral-sh/uv:python3.13-bookworm - when: - event: pull_request - branch: - - main - - dev - path: - - "src/**" - - "tests/**" - environment: - WRENN_API_KEY: - from_secret: WRENN_API_KEY - commands: - - uv sync --dev - - uv run pytest -m integration -v diff --git a/docs/reference.md b/docs/reference.md index 9a406df..49870ff 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -709,7 +709,8 @@ Connect to a running background process and stream its output. #### stream ```python -def stream(cmd: str, args: list[str] | None = None) -> Iterator[StreamEvent] +def stream(cmd: str, + args: builtins.list[str] | None = None) -> Iterator[StreamEvent] ``` Execute a command via WebSocket, streaming output as events. @@ -836,8 +837,9 @@ Connect to a running background process and stream its output. #### stream ```python -async def stream(cmd: str, - args: list[str] | None = None) -> AsyncIterator[StreamEvent] +async def stream( + cmd: str, + args: builtins.list[str] | None = None) -> AsyncIterator[StreamEvent] ``` Execute a command via WebSocket, streaming output as events. @@ -1271,407 +1273,28 @@ in memory. # wrenn.code\_interpreter.models - - -## ExecutionError Objects - -```python -@dataclass -class ExecutionError() -``` - -Error raised during code execution. - -**Attributes**: - -- `name` - Exception class name (e.g. ``"NameError"``). -- `value` - Exception message. -- `traceback` - Full traceback string. - - - -## Logs Objects - -```python -@dataclass -class Logs() -``` - -Captured stdout/stderr streams. - -Each element in the list is one chunk of text as it arrived from -the kernel. - - - -## Result Objects - -```python -@dataclass -class Result() -``` - -A single rich output from code execution. - -Jupyter cells can produce multiple outputs — one ``execute_result`` -(the expression value) and zero or more ``display_data`` messages -(from ``plt.show()``, ``display()``, etc.). Each becomes a -``Result``. - -Known MIME types are unpacked into named attributes; anything else -lands in :pyattr:`extra`. - - - -#### text - -``text/plain`` representation. - - - -#### html - -``text/html`` representation. - - - -#### markdown - -``text/markdown`` representation. - - - -#### svg - -``image/svg+xml`` representation. - - - -#### png - -``image/png`` — base64-encoded. - - - -#### jpeg - -``image/jpeg`` — base64-encoded. - - - -#### pdf - -``application/pdf`` — base64-encoded. - - - -#### latex - -``text/latex`` representation. - - - -#### json - -``application/json`` representation. - - - -#### javascript - -``application/javascript`` representation. - - - -#### extra - -MIME types not covered by the named fields above. - - - -#### is\_main\_result - -``True`` when this came from an ``execute_result`` message -(i.e. the value of the last expression in the cell). ``False`` -for ``display_data`` outputs. - - - -#### from\_bundle - -```python -@classmethod -def from_bundle(cls, - bundle: dict[str, str], - *, - is_main_result: bool = False) -> Result -``` - -Build a ``Result`` from a Jupyter MIME bundle dict. - - - -#### formats - -```python -def formats() -> list[str] -``` - -Return names of non-``None`` MIME-type fields. - - - -## Execution Objects - -```python -@dataclass -class Execution() -``` - -Complete result of a ``run_code`` call. - -**Attributes**: - -- `results` - All rich outputs produced by the cell — charts, tables, - images, expression values, etc. -- `logs` - Captured stdout/stderr text. -- `error` - Populated when the cell raised an exception. -- `execution_count` - Jupyter execution counter (the ``[N]`` number). - - - -#### text - -```python -@property -def text() -> str | None -``` - -Convenience — ``text/plain`` of the main ``execute_result``, -or ``None`` if the cell had no expression value. +Deprecated — use :mod:`wrenn.code_runner.models`. # wrenn.code\_interpreter.async\_capsule - - -## AsyncCapsule Objects - -```python -class AsyncCapsule(BaseAsyncCapsule) -``` - -Async code interpreter capsule with ``run_code`` support. - -Uses ``code-runner-beta`` template by default:: - -from wrenn.code_interpreter import AsyncCapsule - -capsule = await AsyncCapsule.create() -result = await capsule.run_code("print('hello')") - - - -#### create - -```python -@classmethod -async def create(cls, - template: str | None = None, - vcpus: int | None = None, - memory_mb: int | None = None, - timeout: int | None = None, - *, - wait: bool = False, - api_key: str | None = None, - base_url: str | None = None) -> AsyncCapsule -``` - -Create a new async code interpreter capsule. - -**Arguments**: - -- `template` _str | None_ - Template to boot from. Defaults to - ``"code-runner-beta"``. -- `vcpus` _int | None_ - Number of virtual CPUs. -- `memory_mb` _int | None_ - Memory in MiB. -- `timeout` _int | None_ - Inactivity TTL in seconds before auto-pause. -- `wait` _bool_ - Await until the capsule reaches ``running`` status. -- `api_key` _str | None_ - Wrenn API key. Falls back to - ``WRENN_API_KEY`` env var. -- `base_url` _str | None_ - API base URL override. - - -**Returns**: - -- `AsyncCapsule` - A new async code interpreter capsule instance. - - - -#### run\_code - -```python -async def run_code( - code: str, - language: str = "python", - timeout: float = 30, - jupyter_timeout: float = 30, - on_result: Callable[[Result], Any] | None = None, - on_stdout: Callable[[str], Any] | None = None, - on_stderr: Callable[[str], Any] | None = None, - on_error: Callable[[ExecutionError], Any] | None = None) -> Execution -``` - -Execute code in a persistent Jupyter kernel (async). - -**Arguments**: - -- `code` - Code string to execute. -- `language` - Execution backend language. Currently only ``"python"``. -- `timeout` - Maximum seconds to wait for execution to complete. -- `jupyter_timeout` - Maximum seconds to wait for Jupyter to become - available. -- `on_result` - Called for each rich output (charts, images, expression - values). -- `on_stdout` - Called for each stdout chunk. -- `on_stderr` - Called for each stderr chunk. -- `on_error` - Called when the cell raises an exception. - - -**Returns**: - - An :class:`Execution` with ``.results``, ``.logs``, ``.error``, - and a convenience ``.text`` property. +Deprecated — use :mod:`wrenn.code_runner.async_capsule`. # wrenn.code\_interpreter +Deprecated alias for :mod:`wrenn.code_runner`. + +Importing from ``wrenn.code_interpreter`` emits a ``FutureWarning``. +Use ``wrenn.code_runner`` instead. + # wrenn.code\_interpreter.capsule - - -## Capsule Objects - -```python -class Capsule(BaseCapsule) -``` - -Code interpreter capsule with ``run_code`` support. - -Uses ``code-runner-beta`` template by default:: - -from wrenn.code_interpreter import Capsule - -capsule = Capsule() -result = capsule.run_code("print('hello')") -print(result.logs.stdout) # ["hello\n"] - - - -#### \_\_init\_\_ - -```python -def __init__(template: str | None = None, - vcpus: int | None = None, - memory_mb: int | None = None, - timeout: int | None = None, - *, - api_key: str | None = None, - base_url: str | None = None, - **kwargs) -> None -``` - -Create a code interpreter capsule. - -**Arguments**: - -- `template` _str | None_ - Template to boot from. Defaults to - ``"code-runner-beta"``. -- `vcpus` _int | None_ - Number of virtual CPUs. -- `memory_mb` _int | None_ - Memory in MiB. -- `timeout` _int | None_ - Inactivity TTL in seconds before auto-pause. -- `api_key` _str | None_ - Wrenn API key. Falls back to - ``WRENN_API_KEY`` env var. -- `base_url` _str | None_ - API base URL override. - - - -#### create - -```python -@classmethod -def create(cls, - template: str | None = None, - vcpus: int | None = None, - memory_mb: int | None = None, - timeout: int | None = None, - *, - wait: bool = False, - api_key: str | None = None, - base_url: str | None = None) -> Capsule -``` - -Create a new code interpreter capsule. - -**Arguments**: - -- `template` _str | None_ - Template to boot from. Defaults to - ``"code-runner-beta"``. -- `vcpus` _int | None_ - Number of virtual CPUs. -- `memory_mb` _int | None_ - Memory in MiB. -- `timeout` _int | None_ - Inactivity TTL in seconds before auto-pause. -- `wait` _bool_ - Block until the capsule reaches ``running`` status. -- `api_key` _str | None_ - Wrenn API key. Falls back to - ``WRENN_API_KEY`` env var. -- `base_url` _str | None_ - API base URL override. - - -**Returns**: - -- `Capsule` - A new code interpreter capsule instance. - - - -#### run\_code - -```python -def run_code( - code: str, - language: str = "python", - timeout: float = 30, - jupyter_timeout: float = 30, - on_result: Callable[[Result], Any] | None = None, - on_stdout: Callable[[str], Any] | None = None, - on_stderr: Callable[[str], Any] | None = None, - on_error: Callable[[ExecutionError], Any] | None = None) -> Execution -``` - -Execute code in a persistent Jupyter kernel. - -Variables, imports, and function definitions survive across calls. - -**Arguments**: - -- `code` - Code string to execute. -- `language` - Execution backend language. Currently only ``"python"``. -- `timeout` - Maximum seconds to wait for execution to complete. -- `jupyter_timeout` - Maximum seconds to wait for Jupyter to become - available. -- `on_result` - Called for each rich output (charts, images, expression - values). -- `on_stdout` - Called for each stdout chunk. -- `on_stderr` - Called for each stderr chunk. -- `on_error` - Called when the cell raises an exception. - - -**Returns**: - - An :class:`Execution` with ``.results``, ``.logs``, ``.error``, - and a convenience ``.text`` property. +Deprecated — use :mod:`wrenn.code_runner.capsule`. @@ -1964,25 +1587,15 @@ inactivity TTL is set. #### wait\_ready ```python -async def wait_ready(timeout: float = 30) -> None +async def wait_ready(timeout: float = _DEFAULT_WAIT_TIMEOUT) -> None ``` -Await until the capsule status is ``running``. - -Polling interval adapts to the current transient status: -0.5 s for starting/resuming, 2 s for pausing, 1 s for stopping. - -**Arguments**: - -- `timeout` _float_ - Maximum seconds to wait. Defaults to ``30``. - +Await until capsule status is ``running``. **Raises**: -- `TimeoutError` - If the capsule does not reach ``running`` state - within ``timeout`` seconds. -- `RuntimeError` - If the capsule enters an error, stopped, or paused - state while waiting. +- `TimeoutError` - If capsule does not reach ``running`` within ``timeout``. +- `RuntimeError` - If capsule enters error/stopped/missing while waiting. @@ -2032,7 +1645,7 @@ List all capsules belonging to the team. ```python @asynccontextmanager async def pty(cmd: str = "/bin/bash", - args: list[str] | None = None, + args: builtins.list[str] | None = None, cols: int = 80, rows: int = 24, envs: dict[str, str] | None = None, @@ -2094,7 +1707,7 @@ Reconnect to an existing PTY session by tag. def get_url(port: int) -> str ``` -Get the proxy URL for a port exposed inside this capsule. +Get the HTTP proxy URL for a port exposed inside this capsule. **Arguments**: @@ -2103,8 +1716,10 @@ Get the proxy URL for a port exposed inside this capsule. **Returns**: -- `str` - A ``wss://`` (or ``ws://``) URL that proxies to the given - port inside the capsule. +- `str` - A ``https://`` (or ``http://``) URL that proxies HTTP + requests to the given port inside the capsule. For raw + WebSocket access, see the lower-level ``_build_proxy_url`` + helper or the ``pty()`` API. @@ -2309,6 +1924,18 @@ Send SIGKILL to the PTY process. # wrenn.models.\_generated + + +## SessionResponse Objects + +```python +class SessionResponse(BaseModel) +``` + +Returned by login, activate, and switch-team. The actual auth credential +is the wrenn_sid cookie set on the response. The body carries identity +data the SPA needs to bootstrap. + ## Peaks Objects @@ -2349,6 +1976,29 @@ class Type2(StrEnum) Host type. Regular hosts are shared; BYOC hosts belong to a team. + + +## Outcome Objects + +```python +class Outcome(StrEnum) +``` + +Present for action events (capsule.* except state.changed, +template.snapshot.*). Absent for host.up/down, capsule.state.changed, +and the connected sentinel. + + + +## SSEEvent Objects + +```python +class SSEEvent(BaseModel) +``` + +Wire format of one SSE message body. The event name (`event:` line) is +the `kind` and the JSON below is the `data:` line. + # wrenn.models @@ -2536,25 +2186,15 @@ inactivity TTL is set. #### wait\_ready ```python -def wait_ready(timeout: float = 30) -> None +def wait_ready(timeout: float = _DEFAULT_WAIT_TIMEOUT) -> None ``` -Block until the capsule status is ``running``. - -Polling interval adapts to the current transient status: -0.5 s for starting/resuming, 2 s for pausing, 1 s for stopping. - -**Arguments**: - -- `timeout` _float_ - Maximum seconds to wait. Defaults to ``30``. - +Block until capsule status is ``running``. **Raises**: -- `TimeoutError` - If the capsule does not reach ``running`` state - within ``timeout`` seconds. -- `RuntimeError` - If the capsule enters an error, stopped, or paused - state while waiting. +- `TimeoutError` - If capsule does not reach ``running`` within ``timeout``. +- `RuntimeError` - If capsule enters error/stopped/missing while waiting. @@ -2604,7 +2244,7 @@ List all capsules belonging to the team. ```python @contextmanager def pty(cmd: str = "/bin/bash", - args: list[str] | None = None, + args: builtins.list[str] | None = None, cols: int = 80, rows: int = 24, envs: dict[str, str] | None = None, @@ -2665,7 +2305,7 @@ Reconnect to an existing PTY session by tag. def get_url(port: int) -> str ``` -Get the proxy URL for a port exposed inside this capsule. +Get the HTTP proxy URL for a port exposed inside this capsule. **Arguments**: @@ -2674,8 +2314,10 @@ Get the proxy URL for a port exposed inside this capsule. **Returns**: -- `str` - A ``wss://`` (or ``ws://``) URL that proxies to the given - port inside the capsule. +- `str` - A ``https://`` (or ``http://``) URL that proxies HTTP + requests to the given port inside the capsule. For raw + WebSocket access, see the lower-level ``_build_proxy_url`` + helper or the ``pty()`` API. @@ -2700,6 +2342,494 @@ Create a snapshot template from this capsule's current state. - `Template` - The created snapshot template. + + +# wrenn.code\_runner.models + + + +## ExecutionError Objects + +```python +@dataclass +class ExecutionError() +``` + +Error raised during code execution. + +**Attributes**: + +- `name` - Exception class name (e.g. ``"NameError"``). +- `value` - Exception message. +- `traceback` - Full traceback string. + + + +## Logs Objects + +```python +@dataclass +class Logs() +``` + +Captured stdout/stderr streams. + +Each element in the list is one chunk of text as it arrived from +the kernel. + + + +## Result Objects + +```python +@dataclass +class Result() +``` + +A single rich output from code execution. + +Jupyter cells can produce multiple outputs — one ``execute_result`` +(the expression value) and zero or more ``display_data`` messages +(from ``plt.show()``, ``display()``, etc.). Each becomes a +``Result``. + +Known MIME types are unpacked into named attributes; anything else +lands in :pyattr:`extra`. + + + +#### text + +``text/plain`` representation. + + + +#### html + +``text/html`` representation. + + + +#### markdown + +``text/markdown`` representation. + + + +#### svg + +``image/svg+xml`` representation. + + + +#### png + +``image/png`` — base64-encoded. + + + +#### jpeg + +``image/jpeg`` — base64-encoded. + + + +#### gif + +``image/gif`` — base64-encoded. + + + +#### pdf + +``application/pdf`` — base64-encoded. + + + +#### latex + +``text/latex`` representation. + + + +#### json + +``application/json`` representation. + + + +#### javascript + +``application/javascript`` representation. + + + +#### plotly + +``application/vnd.plotly.v1+json`` representation. + + + +#### extra + +MIME types not covered by the named fields above. + + + +#### is\_main\_result + +``True`` when this came from an ``execute_result`` message +(i.e. the value of the last expression in the cell). ``False`` +for ``display_data`` outputs. + + + +#### from\_bundle + +```python +@classmethod +def from_bundle(cls, + bundle: dict[str, str], + *, + is_main_result: bool = False) -> Result +``` + +Build a ``Result`` from a Jupyter MIME bundle dict. + + + +#### formats + +```python +def formats() -> list[str] +``` + +Return names of non-``None`` MIME-type fields. + + + +## Execution Objects + +```python +@dataclass +class Execution() +``` + +Complete result of a ``run_code`` call. + +**Attributes**: + +- `results` - All rich outputs produced by the cell — charts, tables, + images, expression values, etc. +- `logs` - Captured stdout/stderr text. +- `error` - Populated when the cell raised an exception. +- `execution_count` - Jupyter execution counter (the ``[N]`` number). + + + +#### timed\_out + +``True`` when execution was cut short by the ``timeout`` parameter +(or by the kernel WebSocket dropping). Pairs with ``error`` of name +``"Timeout"`` or ``"Disconnected"``. + + + +#### text + +```python +@property +def text() -> str | None +``` + +Convenience — ``text/plain`` of the main ``execute_result``, +or ``None`` if the cell had no expression value. + + + +# wrenn.code\_runner.async\_capsule + + + +## AsyncCapsule Objects + +```python +class AsyncCapsule(BaseAsyncCapsule) +``` + +Async code runner capsule with ``run_code`` support. + +Uses ``code-runner-beta`` template and the ``wrenn`` Jupyter +kernelspec by default:: + +from wrenn.code_runner import AsyncCapsule + +capsule = await AsyncCapsule.create() +result = await capsule.run_code("print('hello')") + + + +#### create + +```python +@classmethod +async def create(cls, + template: str | None = None, + vcpus: int | None = None, + memory_mb: int | None = None, + timeout: int | None = None, + *, + kernel: str | None = None, + wait: bool = False, + api_key: str | None = None, + base_url: str | None = None) -> AsyncCapsule +``` + +Create a new async code runner capsule. + +**Arguments**: + +- `template` _str | None_ - Template to boot from. Defaults to + ``"code-runner-beta"``. +- `vcpus` _int | None_ - Number of virtual CPUs. +- `memory_mb` _int | None_ - Memory in MiB. +- `timeout` _int | None_ - Inactivity TTL in seconds before auto-pause. +- `kernel` _str | None_ - Jupyter kernelspec name. Defaults to + ``"wrenn"``. +- `wait` _bool_ - Await until the capsule reaches ``running`` status. +- `api_key` _str | None_ - Wrenn API key. Falls back to + ``WRENN_API_KEY`` env var. +- `base_url` _str | None_ - API base URL override. + + +**Returns**: + +- `AsyncCapsule` - A new async code runner capsule instance. + + + +#### run\_code + +```python +async def run_code( + code: str, + language: str = "python", + timeout: float = 30, + jupyter_timeout: float = 30, + on_result: Callable[[Result], Any] | None = None, + on_stdout: Callable[[str], Any] | None = None, + on_stderr: Callable[[str], Any] | None = None, + on_error: Callable[[ExecutionError], Any] | None = None) -> Execution +``` + +Execute code in a persistent Jupyter kernel (async). + +**Arguments**: + +- `code` - Code string to execute. +- `language` - Execution backend language. Currently only ``"python"``. +- `timeout` - Maximum seconds to wait for execution to complete. +- `jupyter_timeout` - Maximum seconds to wait for Jupyter to become + available. +- `on_result` - Called for each rich output (charts, images, expression + values). +- `on_stdout` - Called for each stdout chunk. +- `on_stderr` - Called for each stderr chunk. +- `on_error` - Called when the cell raises an exception. + + +**Returns**: + + An :class:`Execution` with ``.results``, ``.logs``, ``.error``, + and a convenience ``.text`` property. + + + +# wrenn.code\_runner + +Code runner — execute code in persistent Jupyter kernels. + +Uses the ``code-runner-beta`` template and the ``wrenn`` Jupyter +kernelspec by default. + +Example:: + +from wrenn.code_runner import Capsule + +with Capsule(wait=True) as capsule: +result = capsule.run_code("print('hello')") +print(result.logs.stdout) + + + +# wrenn.code\_runner.capsule + + + +## Capsule Objects + +```python +class Capsule(BaseCapsule) +``` + +Code runner capsule with ``run_code`` support. + +Uses ``code-runner-beta`` template and the ``wrenn`` Jupyter +kernelspec by default:: + +from wrenn.code_runner import Capsule + +capsule = Capsule() +result = capsule.run_code("print('hello')") +print(result.logs.stdout) # ["hello\n"] + + + +#### \_\_init\_\_ + +```python +def __init__(template: str | None = None, + vcpus: int | None = None, + memory_mb: int | None = None, + timeout: int | None = None, + *, + kernel: str | None = None, + api_key: str | None = None, + base_url: str | None = None, + **kwargs) -> None +``` + +Create a code runner capsule. + +**Arguments**: + +- `template` _str | None_ - Template to boot from. Defaults to + ``"code-runner-beta"``. +- `vcpus` _int | None_ - Number of virtual CPUs. +- `memory_mb` _int | None_ - Memory in MiB. +- `timeout` _int | None_ - Inactivity TTL in seconds before auto-pause. +- `kernel` _str | None_ - Jupyter kernelspec name. Defaults to + ``"wrenn"``. +- `api_key` _str | None_ - Wrenn API key. Falls back to + ``WRENN_API_KEY`` env var. +- `base_url` _str | None_ - API base URL override. + + + +#### create + +```python +@classmethod +def create(cls, + template: str | None = None, + vcpus: int | None = None, + memory_mb: int | None = None, + timeout: int | None = None, + *, + kernel: str | None = None, + wait: bool = False, + api_key: str | None = None, + base_url: str | None = None) -> Capsule +``` + +Create a new code runner capsule. + +**Arguments**: + +- `template` _str | None_ - Template to boot from. Defaults to + ``"code-runner-beta"``. +- `vcpus` _int | None_ - Number of virtual CPUs. +- `memory_mb` _int | None_ - Memory in MiB. +- `timeout` _int | None_ - Inactivity TTL in seconds before auto-pause. +- `kernel` _str | None_ - Jupyter kernelspec name. Defaults to + ``"wrenn"``. +- `wait` _bool_ - Block until the capsule reaches ``running`` status. +- `api_key` _str | None_ - Wrenn API key. Falls back to + ``WRENN_API_KEY`` env var. +- `base_url` _str | None_ - API base URL override. + + +**Returns**: + +- `Capsule` - A new code runner capsule instance. + + + +#### run\_code + +```python +def run_code( + code: str, + language: str = "python", + timeout: float = 30, + jupyter_timeout: float = 30, + on_result: Callable[[Result], Any] | None = None, + on_stdout: Callable[[str], Any] | None = None, + on_stderr: Callable[[str], Any] | None = None, + on_error: Callable[[ExecutionError], Any] | None = None) -> Execution +``` + +Execute code in a persistent Jupyter kernel. + +Variables, imports, and function definitions survive across calls. + +**Arguments**: + +- `code` - Code string to execute. +- `language` - Execution backend language. Currently only ``"python"`` + is supported; passing anything else raises ``ValueError``. + To target a non-Python kernel, set ``kernel=`` on the + capsule constructor. +- `timeout` - Maximum seconds to wait for execution to complete. +- `jupyter_timeout` - Maximum seconds to wait for Jupyter to become + available. +- `on_result` - Called for each rich output (charts, images, expression + values). +- `on_stdout` - Called for each stdout chunk. +- `on_stderr` - Called for each stderr chunk. +- `on_error` - Called when the cell raises an exception. + + +**Returns**: + + An :class:`Execution` with ``.results``, ``.logs``, ``.error``, + and a convenience ``.text`` property. + + + +# wrenn.code\_runner.\_protocol + +Shared Jupyter protocol helpers used by both sync and async capsules. + +Pure functions only — no I/O, no sync/async coupling. + + + +#### build\_execute\_request + +```python +def build_execute_request(code: str) -> dict +``` + +Build a Jupyter ``execute_request`` message envelope. + +**Returns**: + +- `dict` - A fully-formed Jupyter shell-channel message ready to be + JSON-serialized over the kernel WebSocket. The caller is + expected to read ``msg["header"]["msg_id"]`` to correlate + responses. + + + +#### build\_ws\_url + +```python +def build_ws_url(base_url: str, capsule_id: str, kernel_id: str) -> str +``` + +Build the Jupyter kernel WebSocket URL for the given capsule. + # wrenn.\_config @@ -3158,16 +3288,6 @@ def build_config_get(key: str, Build ``git config --get`` arguments. - - -#### build\_has\_upstream - -```python -def build_has_upstream() -> list[str] -``` - -Build arguments to check if current branch has upstream tracking. - #### parse\_status diff --git a/src/wrenn/_git/__init__.py b/src/wrenn/_git/__init__.py index fa59564..05d2722 100644 --- a/src/wrenn/_git/__init__.py +++ b/src/wrenn/_git/__init__.py @@ -153,6 +153,20 @@ class Git: timeout=timeout, ) + def _run_op( + self, + argv: list[str], + *, + op: str, + cwd: str | None = None, + envs: dict[str, str] | None = None, + timeout: int | None = 30, + ) -> CommandResult: + """``_run`` + :func:`_check_result` in one call. Raises on failure.""" + result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) + _check_result(result, op=op) + return result + # ── Repository setup ─────────────────────────────────────── def clone( @@ -203,8 +217,7 @@ class Git: clone_url = embed_credentials(url, username, password) argv = build_clone(clone_url, dest, branch=branch, depth=depth) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="clone") + result = self._run_op(argv, op="clone", cwd=cwd, envs=envs, timeout=timeout) if username and password and not dangerously_store_credentials: sanitized = strip_credentials(clone_url) @@ -248,8 +261,7 @@ class Git: GitCommandError: If init failed. """ argv = build_init(path, bare=bare, initial_branch=initial_branch) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="init") + result = self._run_op(argv, op="init", cwd=cwd, envs=envs, timeout=timeout) return result # ── Staging and committing ───────────────────────────────── @@ -280,8 +292,7 @@ class Git: GitCommandError: If add failed. """ argv = build_add(paths, all=all) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="add") + result = self._run_op(argv, op="add", cwd=cwd, envs=envs, timeout=timeout) return result def commit( @@ -318,8 +329,7 @@ class Git: author_name=author_name, author_email=author_email, ) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="commit") + result = self._run_op(argv, op="commit", cwd=cwd, envs=envs, timeout=timeout) return result # ── Remote sync ──────────────────────────────────────────── @@ -375,8 +385,7 @@ class Git: ) argv = build_push(remote, branch, force=force, set_upstream=set_upstream) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="push") + result = self._run_op(argv, op="push", cwd=cwd, envs=envs, timeout=timeout) return result def pull( @@ -430,8 +439,7 @@ class Git: ) argv = build_pull(remote, branch, rebase=rebase, ff_only=ff_only) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="pull") + result = self._run_op(argv, op="pull", cwd=cwd, envs=envs, timeout=timeout) return result # ── Status and branches ──────────────────────────────────── @@ -456,8 +464,9 @@ class Git: Raises: GitCommandError: If the command failed. """ - result = self._run(build_status(), cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="status") + result = self._run_op( + build_status(), op="status", cwd=cwd, envs=envs, timeout=timeout + ) return parse_status(result.stdout) def branches( @@ -480,8 +489,9 @@ class Git: Raises: GitCommandError: If the command failed. """ - result = self._run(build_branches(), cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="branches") + result = self._run_op( + build_branches(), op="branches", cwd=cwd, envs=envs, timeout=timeout + ) return parse_branches(result.stdout) def create_branch( @@ -509,8 +519,9 @@ class Git: GitCommandError: If the command failed. """ argv = build_create_branch(name, start_point=start_point) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="create_branch") + result = self._run_op( + argv, op="create_branch", cwd=cwd, envs=envs, timeout=timeout + ) return result def checkout_branch( @@ -536,8 +547,9 @@ class Git: GitCommandError: If the command failed. """ argv = build_checkout(name) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="checkout_branch") + result = self._run_op( + argv, op="checkout_branch", cwd=cwd, envs=envs, timeout=timeout + ) return result def delete_branch( @@ -565,8 +577,9 @@ class Git: GitCommandError: If the command failed. """ argv = build_delete_branch(name, force=force) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="delete_branch") + result = self._run_op( + argv, op="delete_branch", cwd=cwd, envs=envs, timeout=timeout + ) return result # ── Remotes ──────────────────────────────────────────────── @@ -598,8 +611,9 @@ class Git: GitCommandError: If the command failed. """ argv = build_remote_add(name, url, fetch=fetch) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="remote_add") + result = self._run_op( + argv, op="remote_add", cwd=cwd, envs=envs, timeout=timeout + ) return result def remote_get( @@ -661,8 +675,7 @@ class Git: GitCommandError: If the command failed. """ argv = build_reset(mode=mode, ref=ref, paths=paths) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="reset") + result = self._run_op(argv, op="reset", cwd=cwd, envs=envs, timeout=timeout) return result def restore( @@ -694,8 +707,7 @@ class Git: GitCommandError: If the command failed. """ argv = build_restore(paths, staged=staged, worktree=worktree, source=source) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="restore") + result = self._run_op(argv, op="restore", cwd=cwd, envs=envs, timeout=timeout) return result # ── Configuration ────────────────────────────────────────── @@ -729,8 +741,9 @@ class Git: GitCommandError: If the command failed. """ argv = build_config_set(key, value, scope=scope, repo_path=cwd) - result = self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="set_config") + result = self._run_op( + argv, op="set_config", cwd=cwd, envs=envs, timeout=timeout + ) return result def get_config( @@ -957,6 +970,20 @@ class AsyncGit: timeout=timeout, ) + async def _run_op( + self, + argv: list[str], + *, + op: str, + cwd: str | None = None, + envs: dict[str, str] | None = None, + timeout: int | None = 30, + ) -> CommandResult: + """``_run`` + :func:`_check_result` in one call. Raises on failure.""" + result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) + _check_result(result, op=op) + return result + # ── Repository setup ─────────────────────────────────────── async def clone( @@ -984,8 +1011,9 @@ class AsyncGit: clone_url = embed_credentials(url, username, password) argv = build_clone(clone_url, dest, branch=branch, depth=depth) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="clone") + result = await self._run_op( + argv, op="clone", cwd=cwd, envs=envs, timeout=timeout + ) if username and password and not dangerously_store_credentials: sanitized = strip_credentials(clone_url) @@ -1014,8 +1042,9 @@ class AsyncGit: ) -> CommandResult: """Initialize a new git repository.""" argv = build_init(path, bare=bare, initial_branch=initial_branch) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="init") + result = await self._run_op( + argv, op="init", cwd=cwd, envs=envs, timeout=timeout + ) return result # ── Staging and committing ───────────────────────────────── @@ -1031,8 +1060,7 @@ class AsyncGit: ) -> CommandResult: """Stage files for commit.""" argv = build_add(paths, all=all) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="add") + result = await self._run_op(argv, op="add", cwd=cwd, envs=envs, timeout=timeout) return result async def commit( @@ -1053,8 +1081,9 @@ class AsyncGit: author_name=author_name, author_email=author_email, ) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="commit") + result = await self._run_op( + argv, op="commit", cwd=cwd, envs=envs, timeout=timeout + ) return result # ── Remote sync ──────────────────────────────────────────── @@ -1095,8 +1124,9 @@ class AsyncGit: ) argv = build_push(remote, branch, force=force, set_upstream=set_upstream) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="push") + result = await self._run_op( + argv, op="push", cwd=cwd, envs=envs, timeout=timeout + ) return result async def pull( @@ -1135,8 +1165,9 @@ class AsyncGit: ) argv = build_pull(remote, branch, rebase=rebase, ff_only=ff_only) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="pull") + result = await self._run_op( + argv, op="pull", cwd=cwd, envs=envs, timeout=timeout + ) return result # ── Status and branches ──────────────────────────────────── @@ -1149,8 +1180,9 @@ class AsyncGit: timeout: int | None = 30, ) -> GitStatus: """Get repository status.""" - result = await self._run(build_status(), cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="status") + result = await self._run_op( + build_status(), op="status", cwd=cwd, envs=envs, timeout=timeout + ) return parse_status(result.stdout) async def branches( @@ -1161,8 +1193,9 @@ class AsyncGit: timeout: int | None = 30, ) -> list[GitBranch]: """List local branches.""" - result = await self._run(build_branches(), cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="branches") + result = await self._run_op( + build_branches(), op="branches", cwd=cwd, envs=envs, timeout=timeout + ) return parse_branches(result.stdout) async def create_branch( @@ -1176,8 +1209,9 @@ class AsyncGit: ) -> CommandResult: """Create and check out a new branch.""" argv = build_create_branch(name, start_point=start_point) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="create_branch") + result = await self._run_op( + argv, op="create_branch", cwd=cwd, envs=envs, timeout=timeout + ) return result async def checkout_branch( @@ -1190,8 +1224,9 @@ class AsyncGit: ) -> CommandResult: """Check out an existing branch.""" argv = build_checkout(name) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="checkout_branch") + result = await self._run_op( + argv, op="checkout_branch", cwd=cwd, envs=envs, timeout=timeout + ) return result async def delete_branch( @@ -1205,8 +1240,9 @@ class AsyncGit: ) -> CommandResult: """Delete a branch.""" argv = build_delete_branch(name, force=force) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="delete_branch") + result = await self._run_op( + argv, op="delete_branch", cwd=cwd, envs=envs, timeout=timeout + ) return result # ── Remotes ──────────────────────────────────────────────── @@ -1223,8 +1259,9 @@ class AsyncGit: ) -> CommandResult: """Add a remote.""" argv = build_remote_add(name, url, fetch=fetch) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="remote_add") + result = await self._run_op( + argv, op="remote_add", cwd=cwd, envs=envs, timeout=timeout + ) return result async def remote_get( @@ -1258,8 +1295,9 @@ class AsyncGit: ) -> CommandResult: """Reset the current HEAD.""" argv = build_reset(mode=mode, ref=ref, paths=paths) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="reset") + result = await self._run_op( + argv, op="reset", cwd=cwd, envs=envs, timeout=timeout + ) return result async def restore( @@ -1275,8 +1313,9 @@ class AsyncGit: ) -> CommandResult: """Restore working-tree files or unstage changes.""" argv = build_restore(paths, staged=staged, worktree=worktree, source=source) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="restore") + result = await self._run_op( + argv, op="restore", cwd=cwd, envs=envs, timeout=timeout + ) return result # ── Configuration ────────────────────────────────────────── @@ -1293,8 +1332,9 @@ class AsyncGit: ) -> CommandResult: """Set a git config value.""" argv = build_config_set(key, value, scope=scope, repo_path=cwd) - result = await self._run(argv, cwd=cwd, envs=envs, timeout=timeout) - _check_result(result, op="set_config") + result = await self._run_op( + argv, op="set_config", cwd=cwd, envs=envs, timeout=timeout + ) return result async def get_config( diff --git a/src/wrenn/_git/_cmd.py b/src/wrenn/_git/_cmd.py index 8e929bf..45bf595 100644 --- a/src/wrenn/_git/_cmd.py +++ b/src/wrenn/_git/_cmd.py @@ -351,11 +351,6 @@ def build_config_get( return args -def build_has_upstream() -> list[str]: - """Build arguments to check if current branch has upstream tracking.""" - return ["git", "rev-parse", "--abbrev-ref", "--symbolic-full-name", "@{u}"] - - # ── Parsers ──────────────────────────────────────────────────────── diff --git a/src/wrenn/async_capsule.py b/src/wrenn/async_capsule.py index 4cf4c96..f091649 100644 --- a/src/wrenn/async_capsule.py +++ b/src/wrenn/async_capsule.py @@ -18,7 +18,7 @@ from wrenn.capsule import ( _RESUME_INTERVAL, _START_INTERVAL, _DualMethod, - _build_proxy_url, + _build_http_proxy_url, ) from wrenn.client import AsyncWrennClient from wrenn.commands import AsyncCommands @@ -423,16 +423,18 @@ class AsyncCapsule: # ── Proxy helpers ─────────────────────────────────────────── def get_url(self, port: int) -> str: - """Get the proxy URL for a port exposed inside this capsule. + """Get the HTTP proxy URL for a port exposed inside this capsule. Args: port (int): Port number to proxy. Returns: - str: A ``wss://`` (or ``ws://``) URL that proxies to the given - port inside the capsule. + str: A ``https://`` (or ``http://``) URL that proxies HTTP + requests to the given port inside the capsule. For raw + WebSocket access, see the lower-level ``_build_proxy_url`` + helper or the ``pty()`` API. """ - return _build_proxy_url(self._client._base_url, self._id, port) + return _build_http_proxy_url(self._client._base_url, self._id, port) # ── Snapshots ─────────────────────────────────────────────── diff --git a/src/wrenn/capsule.py b/src/wrenn/capsule.py index f533205..a5545d5 100644 --- a/src/wrenn/capsule.py +++ b/src/wrenn/capsule.py @@ -21,6 +21,7 @@ from wrenn.pty import PtySession def _build_proxy_url(base_url: str, capsule_id: str | None, port: int) -> str: + """Build the WebSocket proxy URL (``ws://`` / ``wss://``).""" parsed = httpx.URL(base_url) host = parsed.host if parsed.port: @@ -29,6 +30,21 @@ def _build_proxy_url(base_url: str, capsule_id: str | None, port: int) -> str: return f"{scheme}://{port}-{capsule_id}.{host}" +def _build_http_proxy_url(base_url: str, capsule_id: str | None, port: int) -> str: + """Build the HTTP proxy URL (``http://`` / ``https://``). + + The capsule's API base URL typically carries an ``/api`` path suffix + (e.g. ``https://app.wrenn.dev/api``). The proxy host is derived from + the URL's host only — any path is discarded. + """ + parsed = httpx.URL(base_url) + host = parsed.host + if parsed.port: + host = f"{host}:{parsed.port}" + scheme = "http" if parsed.scheme in ("http", "ws") else "https" + return f"{scheme}://{port}-{capsule_id}.{host}" + + _RESUME_INTERVAL = 0.5 _DESTROY_INTERVAL = 0.5 _PAUSE_INTERVAL = 2.0 @@ -499,16 +515,18 @@ class Capsule: # ── Proxy helpers ─────────────────────────────────────────── def get_url(self, port: int) -> str: - """Get the proxy URL for a port exposed inside this capsule. + """Get the HTTP proxy URL for a port exposed inside this capsule. Args: port (int): Port number to proxy. Returns: - str: A ``wss://`` (or ``ws://``) URL that proxies to the given - port inside the capsule. + str: A ``https://`` (or ``http://``) URL that proxies HTTP + requests to the given port inside the capsule. For raw + WebSocket access, see the lower-level ``_build_proxy_url`` + helper or the ``pty()`` API. """ - return _build_proxy_url(self._client._base_url, self._id, port) + return _build_http_proxy_url(self._client._base_url, self._id, port) # ── Snapshots ─────────────────────────────────────────────── diff --git a/src/wrenn/code_runner/_protocol.py b/src/wrenn/code_runner/_protocol.py new file mode 100644 index 0000000..42b5978 --- /dev/null +++ b/src/wrenn/code_runner/_protocol.py @@ -0,0 +1,51 @@ +"""Shared Jupyter protocol helpers used by both sync and async capsules. + +Pure functions only — no I/O, no sync/async coupling. +""" + +from __future__ import annotations + +import time +import uuid + +from wrenn.capsule import _build_proxy_url + + +def build_execute_request(code: str) -> dict: + """Build a Jupyter ``execute_request`` message envelope. + + Returns: + dict: A fully-formed Jupyter shell-channel message ready to be + JSON-serialized over the kernel WebSocket. The caller is + expected to read ``msg["header"]["msg_id"]`` to correlate + responses. + """ + msg_id = str(uuid.uuid4()) + return { + "header": { + "msg_id": msg_id, + "msg_type": "execute_request", + "username": "wrenn-sdk", + "session": str(uuid.uuid4()), + "date": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()), + "version": "5.3", + }, + "parent_header": {}, + "metadata": {}, + "content": { + "code": code, + "silent": False, + "store_history": True, + "user_expressions": {}, + "allow_stdin": False, + "stop_on_error": True, + }, + "buffers": [], + "channel": "shell", + } + + +def build_ws_url(base_url: str, capsule_id: str, kernel_id: str) -> str: + """Build the Jupyter kernel WebSocket URL for the given capsule.""" + proxy = _build_proxy_url(base_url, capsule_id, 8888) + return f"{proxy}/api/kernels/{kernel_id}/channels" diff --git a/src/wrenn/code_runner/async_capsule.py b/src/wrenn/code_runner/async_capsule.py index b8607b3..e11dca0 100644 --- a/src/wrenn/code_runner/async_capsule.py +++ b/src/wrenn/code_runner/async_capsule.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import json import time -import uuid from collections.abc import Callable from typing import Any @@ -11,8 +10,9 @@ import httpx import httpx_ws from wrenn.async_capsule import AsyncCapsule as BaseAsyncCapsule -from wrenn.capsule import _build_proxy_url +from wrenn.capsule import _build_http_proxy_url from wrenn.client import AsyncWrennClient +from wrenn.code_runner._protocol import build_execute_request, build_ws_url from wrenn.code_runner.capsule import DEFAULT_KERNEL, DEFAULT_TEMPLATE from wrenn.code_runner.models import ( Execution, @@ -110,11 +110,7 @@ class AsyncCapsule(BaseAsyncCapsule): def _get_proxy_client(self) -> httpx.AsyncClient: if self._proxy_client is None: - url = ( - _build_proxy_url(self._client._base_url, self._id, 8888) - .replace("ws://", "http://") - .replace("wss://", "https://") - ) + url = _build_http_proxy_url(self._client._base_url, self._id, 8888) self._proxy_client = httpx.AsyncClient( base_url=url, headers={"X-API-Key": self._client._api_key}, @@ -164,36 +160,6 @@ class AsyncCapsule(BaseAsyncCapsule): f"Jupyter not available within {jupyter_timeout}s: {last_exc}" ) - def _jupyter_ws_url(self, kernel_id: str) -> str: - proxy = _build_proxy_url(self._client._base_url, self._id, 8888) - return f"{proxy}/api/kernels/{kernel_id}/channels" - - @staticmethod - def _jupyter_execute_request(code: str) -> dict: - msg_id = str(uuid.uuid4()) - return { - "header": { - "msg_id": msg_id, - "msg_type": "execute_request", - "username": "wrenn-sdk", - "session": str(uuid.uuid4()), - "date": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()), - "version": "5.3", - }, - "parent_header": {}, - "metadata": {}, - "content": { - "code": code, - "silent": False, - "store_history": True, - "user_expressions": {}, - "allow_stdin": False, - "stop_on_error": True, - }, - "buffers": [], - "channel": "shell", - } - async def run_code( self, code: str, @@ -230,24 +196,42 @@ class AsyncCapsule(BaseAsyncCapsule): "non-Python kernelspec." ) kernel_id = await self._ensure_kernel(jupyter_timeout=jupyter_timeout) - ws_url = self._jupyter_ws_url(kernel_id) + ws_url = build_ws_url(self._client._base_url, self._id, kernel_id) - msg = self._jupyter_execute_request(code) + msg = build_execute_request(code) msg_id = msg["header"]["msg_id"] execution = Execution() deadline = time.monotonic() + timeout headers = {"X-API-Key": self._client._api_key} + saw_idle = False + + def _emit_error(err: ExecutionError) -> None: + execution.error = err + if on_error is not None: + on_error(err) async with httpx_ws.aconnect_ws(ws_url, headers=headers) as ws: # type: httpx_ws.AsyncWebSocketSession await ws.send_text(json.dumps(msg)) - while time.monotonic() < deadline: + while True: time_left = deadline - time.monotonic() if time_left <= 0: break try: data = await asyncio.wait_for(ws.receive_json(), timeout=time_left) - except Exception: + except (asyncio.TimeoutError, TimeoutError): + break + except ( + httpx_ws.WebSocketDisconnect, + httpx_ws.WebSocketNetworkError, + ) as exc: + execution.timed_out = True + _emit_error( + ExecutionError( + name="Disconnected", + value=f"kernel WebSocket closed: {exc}", + ) + ) break if not data: break @@ -280,17 +264,26 @@ class AsyncCapsule(BaseAsyncCapsule): if on_result is not None: on_result(result) elif msg_type == "error": - err = ExecutionError( - name=content.get("ename", ""), - value=content.get("evalue", ""), - traceback="\n".join(content.get("traceback", [])), + _emit_error( + ExecutionError( + name=content.get("ename", ""), + value=content.get("evalue", ""), + traceback="\n".join(content.get("traceback", [])), + ) ) - execution.error = err - if on_error is not None: - on_error(err) elif msg_type == "status" and content.get("execution_state") == "idle": + saw_idle = True break + if not saw_idle and execution.error is None: + execution.timed_out = True + _emit_error( + ExecutionError( + name="Timeout", + value=f"run_code exceeded {timeout}s", + ) + ) + return execution async def __aexit__(self, *args) -> None: diff --git a/src/wrenn/code_runner/capsule.py b/src/wrenn/code_runner/capsule.py index cac94b0..782e812 100644 --- a/src/wrenn/code_runner/capsule.py +++ b/src/wrenn/code_runner/capsule.py @@ -2,7 +2,6 @@ from __future__ import annotations import json import time -import uuid from collections.abc import Callable from typing import Any @@ -10,7 +9,8 @@ import httpx import httpx_ws from wrenn.capsule import Capsule as BaseCapsule -from wrenn.capsule import _build_proxy_url +from wrenn.capsule import _build_http_proxy_url +from wrenn.code_runner._protocol import build_execute_request, build_ws_url from wrenn.code_runner.models import ( Execution, ExecutionError, @@ -138,11 +138,7 @@ class Capsule(BaseCapsule): def _get_proxy_client(self) -> httpx.Client: if self._proxy_client is None: - url = ( - _build_proxy_url(self._client._base_url, self._id, 8888) - .replace("ws://", "http://") - .replace("wss://", "https://") - ) + url = _build_http_proxy_url(self._client._base_url, self._id, 8888) self._proxy_client = httpx.Client( base_url=url, headers={"X-API-Key": self._client._api_key}, @@ -194,36 +190,6 @@ class Capsule(BaseCapsule): f"Jupyter not available within {jupyter_timeout}s: {last_exc}" ) - def _jupyter_ws_url(self, kernel_id: str) -> str: - proxy = _build_proxy_url(self._client._base_url, self._id, 8888) - return f"{proxy}/api/kernels/{kernel_id}/channels" - - @staticmethod - def _jupyter_execute_request(code: str) -> dict: - msg_id = str(uuid.uuid4()) - return { - "header": { - "msg_id": msg_id, - "msg_type": "execute_request", - "username": "wrenn-sdk", - "session": str(uuid.uuid4()), - "date": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()), - "version": "5.3", - }, - "parent_header": {}, - "metadata": {}, - "content": { - "code": code, - "silent": False, - "store_history": True, - "user_expressions": {}, - "allow_stdin": False, - "stop_on_error": True, - }, - "buffers": [], - "channel": "shell", - } - def run_code( self, code: str, @@ -265,24 +231,42 @@ class Capsule(BaseCapsule): "non-Python kernelspec." ) kernel_id = self._ensure_kernel(jupyter_timeout=jupyter_timeout) - ws_url = self._jupyter_ws_url(kernel_id) + ws_url = build_ws_url(self._client._base_url, self._id, kernel_id) - msg = self._jupyter_execute_request(code) + msg = build_execute_request(code) msg_id = msg["header"]["msg_id"] execution = Execution() deadline = time.monotonic() + timeout headers = {"X-API-Key": self._client._api_key} + saw_idle = False + + def _emit_error(err: ExecutionError) -> None: + execution.error = err + if on_error is not None: + on_error(err) with httpx_ws.connect_ws(ws_url, headers=headers) as ws: # type: httpx_ws.WebSocketSession ws.send_text(json.dumps(msg)) - while time.monotonic() < deadline: + while True: time_left = deadline - time.monotonic() if time_left <= 0: break try: data = ws.receive_json(timeout=time_left) - except Exception: + except TimeoutError: + break + except ( + httpx_ws.WebSocketDisconnect, + httpx_ws.WebSocketNetworkError, + ) as exc: + execution.timed_out = True + _emit_error( + ExecutionError( + name="Disconnected", + value=f"kernel WebSocket closed: {exc}", + ) + ) break if not data: break @@ -315,17 +299,26 @@ class Capsule(BaseCapsule): if on_result is not None: on_result(result) elif msg_type == "error": - err = ExecutionError( - name=content.get("ename", ""), - value=content.get("evalue", ""), - traceback="\n".join(content.get("traceback", [])), + _emit_error( + ExecutionError( + name=content.get("ename", ""), + value=content.get("evalue", ""), + traceback="\n".join(content.get("traceback", [])), + ) ) - execution.error = err - if on_error is not None: - on_error(err) elif msg_type == "status" and content.get("execution_state") == "idle": + saw_idle = True break + if not saw_idle and execution.error is None: + execution.timed_out = True + _emit_error( + ExecutionError( + name="Timeout", + value=f"run_code exceeded {timeout}s", + ) + ) + return execution def __exit__(self, *args) -> None: diff --git a/src/wrenn/code_runner/models.py b/src/wrenn/code_runner/models.py index 39a1e64..42d40ca 100644 --- a/src/wrenn/code_runner/models.py +++ b/src/wrenn/code_runner/models.py @@ -9,10 +9,12 @@ _MIME_MAP: dict[str, str] = { "image/svg+xml": "svg", "image/png": "png", "image/jpeg": "jpeg", + "image/gif": "gif", "application/pdf": "pdf", "text/latex": "latex", "application/json": "json", "application/javascript": "javascript", + "application/vnd.plotly.v1+json": "plotly", } @@ -69,6 +71,8 @@ class Result: """``image/png`` — base64-encoded.""" jpeg: str | None = None """``image/jpeg`` — base64-encoded.""" + gif: str | None = None + """``image/gif`` — base64-encoded.""" pdf: str | None = None """``application/pdf`` — base64-encoded.""" latex: str | None = None @@ -77,6 +81,8 @@ class Result: """``application/json`` representation.""" javascript: str | None = None """``application/javascript`` representation.""" + plotly: dict | None = None + """``application/vnd.plotly.v1+json`` representation.""" extra: dict[str, str] | None = None """MIME types not covered by the named fields above.""" @@ -104,21 +110,9 @@ class Result: def formats(self) -> list[str]: """Return names of non-``None`` MIME-type fields.""" - out: list[str] = [] - for attr in ( - "text", - "html", - "markdown", - "svg", - "png", - "jpeg", - "pdf", - "latex", - "json", - "javascript", - ): - if getattr(self, attr) is not None: - out.append(attr) + out: list[str] = [ + attr for attr in _MIME_MAP.values() if getattr(self, attr) is not None + ] if self.extra: out.extend(self.extra) return out @@ -140,6 +134,10 @@ class Execution: logs: Logs = field(default_factory=Logs) error: ExecutionError | None = None execution_count: int | None = None + timed_out: bool = False + """``True`` when execution was cut short by the ``timeout`` parameter + (or by the kernel WebSocket dropping). Pairs with ``error`` of name + ``"Timeout"`` or ``"Disconnected"``.""" @property def text(self) -> str | None: diff --git a/src/wrenn/files.py b/src/wrenn/files.py index 5a99289..291ff8b 100644 --- a/src/wrenn/files.py +++ b/src/wrenn/files.py @@ -199,7 +199,8 @@ class Files: f"/v1/capsules/{self._capsule_id}/files/stream/write", content=_multipart(), headers={ - "Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}" + "Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}", + "Transfer-Encoding": "chunked", }, ) _raise_for_status(resp) @@ -392,7 +393,8 @@ class AsyncFiles: f"/v1/capsules/{self._capsule_id}/files/stream/write", content=_multipart(), headers={ - "Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}" + "Content-Type": f"multipart/form-data; boundary={boundary.decode('utf-8')}", + "Transfer-Encoding": "chunked", }, ) _raise_for_status(resp) diff --git a/src/wrenn/pty.py b/src/wrenn/pty.py index 63dd26f..0b7ff77 100644 --- a/src/wrenn/pty.py +++ b/src/wrenn/pty.py @@ -53,7 +53,16 @@ def _parse_pty_event(raw: dict[str, Any]) -> PtyEvent: ) if msg_type == "ping": return PtyEvent(type=PtyEventType.ping) - return PtyEvent(type=PtyEventType(msg_type) if msg_type else PtyEventType.ping) + if not msg_type: + return PtyEvent(type=PtyEventType.ping) + try: + return PtyEvent(type=PtyEventType(msg_type)) + except ValueError: + return PtyEvent( + type=PtyEventType.error, + data=f"unknown msg_type: {msg_type!r}", + fatal=False, + ) class PtySession: diff --git a/tests/test_capsule_features.py b/tests/test_capsule_features.py index 7cfe624..186d247 100644 --- a/tests/test_capsule_features.py +++ b/tests/test_capsule_features.py @@ -1,12 +1,14 @@ from __future__ import annotations import httpx +import pytest import respx -from wrenn.capsule import Capsule, _build_proxy_url +from wrenn.capsule import Capsule, _build_http_proxy_url, _build_proxy_url from wrenn.code_runner.models import Execution, ExecutionError, Logs, Result BASE = "https://app.wrenn.dev/api" +API_KEY = "wrn_test1234567890abcdef12345678" class TestBuildProxyUrl: @@ -27,6 +29,23 @@ class TestBuildProxyUrl: assert url == "ws://5000-sb-2.192.168.1.1" +class TestBuildHttpProxyUrl: + """``get_url`` returns an HTTP(S) URL; ``/api`` path on the base URL is + discarded — only the host is used to build the proxy subdomain.""" + + def test_https_production_strips_api_path(self): + url = _build_http_proxy_url("https://app.wrenn.dev/api", "cl-abc", 8080) + assert url == "https://8080-cl-abc.app.wrenn.dev" + + def test_http_localhost_preserves_port(self): + url = _build_http_proxy_url("http://localhost:8080/api", "cl-abc", 3000) + assert url == "http://3000-cl-abc.localhost:8080" + + def test_https_custom_port(self): + url = _build_http_proxy_url("https://api.example.com:9443", "sb-1", 80) + assert url == "https://80-sb-1.api.example.com:9443" + + class TestCapsuleCreate: @respx.mock def test_capsule_constructor_creates(self): @@ -194,6 +213,189 @@ class TestExecutionModels: assert "".join(logs.stderr) == "warn\n" +class TestGetUrlPublic: + """``Capsule.get_url`` returns the HTTP proxy URL.""" + + @respx.mock + def test_sync_get_url_default_base(self): + respx.post(f"{BASE}/v1/capsules").respond( + 202, json={"id": "cl-99", "status": "starting"} + ) + cap = Capsule(api_key=API_KEY, base_url=BASE) + assert cap.get_url(8080) == "https://8080-cl-99.app.wrenn.dev" + + @respx.mock + def test_sync_get_url_localhost(self): + local_base = "http://localhost:8080/api" + respx.post(f"{local_base}/v1/capsules").respond( + 202, json={"id": "cl-42", "status": "starting"} + ) + cap = Capsule(api_key=API_KEY, base_url=local_base) + assert cap.get_url(3000) == "http://3000-cl-42.localhost:8080" + + @pytest.mark.asyncio + @respx.mock + async def test_async_get_url(self): + from wrenn.async_capsule import AsyncCapsule + + respx.post(f"{BASE}/v1/capsules").respond( + 202, json={"id": "cl-async", "status": "starting"} + ) + cap = await AsyncCapsule.create(api_key=API_KEY, base_url=BASE) + assert cap.get_url(5000) == "https://5000-cl-async.app.wrenn.dev" + await cap._client.aclose() + + +class TestPtyConnect: + """``pty_connect`` reconnects to an existing PTY session by tag.""" + + def _capsule(self): + with respx.mock: + respx.post(f"{BASE}/v1/capsules").respond( + 202, json={"id": "cl-1", "status": "starting"} + ) + return Capsule(api_key=API_KEY, base_url=BASE) + + def test_sync_pty_connect_sends_connect_frame(self): + from unittest.mock import MagicMock, patch + + cap = self._capsule() + ws = MagicMock() + ctx = MagicMock() + ctx.__enter__.return_value = ws + ctx.__exit__.return_value = False + + with patch("wrenn.capsule.httpx_ws.connect_ws", return_value=ctx): + with cap.pty_connect("tag-xyz") as session: + assert session is not None + # First send_text call must be a ``connect`` frame with the tag. + import json as _json + + sent = ws.send_text.call_args_list[0].args[0] + payload = _json.loads(sent) + assert payload == {"type": "connect", "tag": "tag-xyz"} + + @pytest.mark.asyncio + @respx.mock + async def test_async_pty_connect_sends_connect_frame(self): + from unittest.mock import AsyncMock, MagicMock, patch + + from wrenn.async_capsule import AsyncCapsule + + respx.post(f"{BASE}/v1/capsules").respond( + 202, json={"id": "cl-1", "status": "starting"} + ) + cap = await AsyncCapsule.create(api_key=API_KEY, base_url=BASE) + ws = MagicMock() + ws.send_text = AsyncMock() + ctx = MagicMock() + ctx.__aenter__ = AsyncMock(return_value=ws) + ctx.__aexit__ = AsyncMock(return_value=False) + + with patch("wrenn.async_capsule.httpx_ws.aconnect_ws", return_value=ctx): + async with cap.pty_connect("tag-async") as session: + assert session is not None + import json as _json + + sent = ws.send_text.call_args_list[0].args[0] + payload = _json.loads(sent) + assert payload == {"type": "connect", "tag": "tag-async"} + await cap._client.aclose() + + +class TestCreateSnapshot: + @respx.mock + def test_sync_create_snapshot_posts_capsule_id(self): + respx.post(f"{BASE}/v1/capsules").respond( + 202, json={"id": "cl-1", "status": "starting"} + ) + snap_route = respx.post(f"{BASE}/v1/snapshots").respond( + 201, + json={"name": "my-snap"}, + ) + cap = Capsule(api_key=API_KEY, base_url=BASE) + tpl = cap.create_snapshot(name="my-snap", overwrite=True) + import json as _json + + req = snap_route.calls[0].request + body = _json.loads(req.content) + assert body["sandbox_id"] == "cl-1" + assert body["name"] == "my-snap" + assert req.url.params["overwrite"] == "true" + assert tpl.name == "my-snap" + + @pytest.mark.asyncio + @respx.mock + async def test_async_create_snapshot(self): + from wrenn.async_capsule import AsyncCapsule + + respx.post(f"{BASE}/v1/capsules").respond( + 202, json={"id": "cl-1", "status": "starting"} + ) + respx.post(f"{BASE}/v1/snapshots").respond( + 201, + json={"name": "auto-named"}, + ) + cap = await AsyncCapsule.create(api_key=API_KEY, base_url=BASE) + tpl = await cap.create_snapshot() + assert tpl.name == "auto-named" + await cap._client.aclose() + + +class TestUploadStreamChunked: + """``upload_stream`` must declare ``Transfer-Encoding: chunked`` and + deliver the multipart body without buffering.""" + + @respx.mock + def test_sync_upload_stream_chunked(self): + respx.post(f"{BASE}/v1/capsules").respond( + 202, json={"id": "cl-1", "status": "starting"} + ) + route = respx.post(f"{BASE}/v1/capsules/cl-1/files/stream/write").respond( + 200, json={} + ) + cap = Capsule(api_key=API_KEY, base_url=BASE) + + def chunks(): + yield b"hello " + yield b"world\n" + + cap.files.upload_stream("/tmp/out.txt", chunks()) + req = route.calls[0].request + assert req.headers["transfer-encoding"] == "chunked" + ct = req.headers["content-type"] + assert ct.startswith("multipart/form-data; boundary=") + body = bytes(req.content) + assert b'name="path"' in body + assert b"/tmp/out.txt" in body + assert b'name="file"' in body + assert b"hello world\n" in body + + @pytest.mark.asyncio + @respx.mock + async def test_async_upload_stream_chunked(self): + from wrenn.async_capsule import AsyncCapsule + + respx.post(f"{BASE}/v1/capsules").respond( + 202, json={"id": "cl-1", "status": "starting"} + ) + route = respx.post(f"{BASE}/v1/capsules/cl-1/files/stream/write").respond( + 200, json={} + ) + cap = await AsyncCapsule.create(api_key=API_KEY, base_url=BASE) + + async def chunks(): + yield b"abc" + yield b"def" + + await cap.files.upload_stream("/tmp/out.bin", chunks()) + req = route.calls[0].request + assert req.headers["transfer-encoding"] == "chunked" + body = bytes(req.content) + assert b"abcdef" in body + await cap._client.aclose() + + class TestDeprecationWarnings: def test_import_sandbox_from_wrenn_warns(self): import sys diff --git a/tests/test_code_runner_unit.py b/tests/test_code_runner_unit.py index c1e3873..94571e0 100644 --- a/tests/test_code_runner_unit.py +++ b/tests/test_code_runner_unit.py @@ -362,12 +362,14 @@ class TestEnsureKernel: c._ensure_kernel(jupyter_timeout=0.01) -# ───────────────────────── _jupyter_execute_request ───────────────────────── +# ───────────────────────── build_execute_request ───────────────────────── class TestJupyterRequest: def test_structure(self): - msg = Capsule._jupyter_execute_request("print(1)") + from wrenn.code_runner._protocol import build_execute_request + + msg = build_execute_request("print(1)") assert msg["channel"] == "shell" assert msg["header"]["msg_type"] == "execute_request" assert msg["content"]["code"] == "print(1)" @@ -379,8 +381,10 @@ class TestJupyterRequest: assert len(msg["header"]["msg_id"]) == 36 def test_unique_msg_id_per_call(self): - a = Capsule._jupyter_execute_request("x") - b = Capsule._jupyter_execute_request("x") + from wrenn.code_runner._protocol import build_execute_request + + a = build_execute_request("x") + b = build_execute_request("x") assert a["header"]["msg_id"] != b["header"]["msg_id"] @@ -397,7 +401,12 @@ def _wrap(msg_type: str, parent_id: str, content: dict) -> dict: class _FakeWS: - """Minimal sync httpx_ws-shaped fake.""" + """Minimal sync httpx_ws-shaped fake. + + If ``frames_factory`` yields an ``Exception`` instance, the fake + raises it instead of returning the value — useful for testing + disconnect / network-error paths. + """ def __init__(self, frames_factory): self._frames_factory = frames_factory @@ -418,9 +427,12 @@ class _FakeWS: def receive_json(self, timeout: float = 0): assert self._iter is not None try: - return next(self._iter) + nxt = next(self._iter) except StopIteration: raise TimeoutError("no more frames") + if isinstance(nxt, BaseException): + raise nxt + return nxt class _FakeAsyncWS: @@ -438,12 +450,15 @@ class _FakeAsyncWS: parent_id = json.loads(s)["header"]["msg_id"] self._iter = iter(self._frames_factory(parent_id)) - async def receive_json(self, timeout: float = 0): + async def receive_json(self): assert self._iter is not None try: - return next(self._iter) + nxt = next(self._iter) except StopIteration: raise TimeoutError("no more frames") + if isinstance(nxt, BaseException): + raise nxt + return nxt class TestRunCode: @@ -630,3 +645,243 @@ class TestAsyncCtorFailureSafe: c = AsyncCapsule.__new__(AsyncCapsule) # __del__ should be safe even with no attrs. c.__del__() + + +# ───────────────────────── run_code error-path regressions (B2) ───────────── + + +class TestRunCodeErrorPaths: + """Sync run_code timeout / disconnect / unexpected-exception behavior.""" + + def _ready(self): + return TestRunCode()._make_ready() + + def test_timeout_when_no_idle_received(self): + c = self._ready() + + def frames(pid): + yield _wrap("stream", pid, {"name": "stdout", "text": "partial\n"}) + # No idle frame; loop exits via StopIteration → TimeoutError. + + errors = [] + with patch( + "wrenn.code_runner.capsule.httpx_ws.connect_ws", + return_value=_FakeWS(frames), + ): + ex = c.run_code("x", on_error=errors.append) + assert ex.timed_out is True + assert ex.error is not None + assert ex.error.name == "Timeout" + assert "exceeded" in ex.error.value + assert ex.logs.stdout == ["partial\n"] + assert len(errors) == 1 + + def test_disconnect_sets_disconnected_error(self): + c = self._ready() + import httpx_ws + + def frames(pid): + yield _wrap("stream", pid, {"name": "stdout", "text": "hi\n"}) + yield httpx_ws.WebSocketDisconnect(code=1000, reason="bye") + + errors = [] + with patch( + "wrenn.code_runner.capsule.httpx_ws.connect_ws", + return_value=_FakeWS(frames), + ): + ex = c.run_code("x", on_error=errors.append) + assert ex.timed_out is True + assert ex.error is not None + assert ex.error.name == "Disconnected" + assert ex.logs.stdout == ["hi\n"] + assert len(errors) == 1 + + def test_unexpected_exception_propagates(self): + c = self._ready() + + def frames(pid): + yield RuntimeError("WS broken in unexpected way") + + with patch( + "wrenn.code_runner.capsule.httpx_ws.connect_ws", + return_value=_FakeWS(frames), + ): + with pytest.raises(RuntimeError, match="WS broken"): + c.run_code("x") + + def test_clean_exit_does_not_set_timed_out(self): + c = self._ready() + + def frames(pid): + yield _wrap("status", pid, {"execution_state": "idle"}) + + with patch( + "wrenn.code_runner.capsule.httpx_ws.connect_ws", + return_value=_FakeWS(frames), + ): + ex = c.run_code("pass") + assert ex.timed_out is False + assert ex.error is None + + +# ───────────────────────── Async run_code parity ────────────────────────── + + +class TestAsyncRunCodeErrorPaths: + def _ready(self): + return TestAsyncRunCode()._make_ready() + + @pytest.mark.asyncio + async def test_async_timeout_when_no_idle(self): + c = self._ready() + + def frames(pid): + yield _wrap("stream", pid, {"name": "stdout", "text": "partial\n"}) + + errors = [] + with patch( + "wrenn.code_runner.async_capsule.httpx_ws.aconnect_ws", + return_value=_FakeAsyncWS(frames), + ): + ex = await c.run_code("x", on_error=errors.append) + assert ex.timed_out is True + assert ex.error is not None + assert ex.error.name == "Timeout" + assert ex.logs.stdout == ["partial\n"] + assert len(errors) == 1 + await c.close() + + @pytest.mark.asyncio + async def test_async_disconnect_sets_disconnected_error(self): + c = self._ready() + import httpx_ws + + def frames(pid): + yield httpx_ws.WebSocketNetworkError("network blip") + + errors = [] + with patch( + "wrenn.code_runner.async_capsule.httpx_ws.aconnect_ws", + return_value=_FakeAsyncWS(frames), + ): + ex = await c.run_code("x", on_error=errors.append) + assert ex.timed_out is True + assert ex.error is not None + assert ex.error.name == "Disconnected" + assert len(errors) == 1 + await c.close() + + @pytest.mark.asyncio + async def test_async_unexpected_exception_propagates(self): + c = self._ready() + + def frames(pid): + yield RuntimeError("unexpected WS death") + + with patch( + "wrenn.code_runner.async_capsule.httpx_ws.aconnect_ws", + return_value=_FakeAsyncWS(frames), + ): + with pytest.raises(RuntimeError, match="unexpected WS"): + await c.run_code("x") + await c.close() + + @pytest.mark.asyncio + async def test_async_unsupported_language_raises(self): + c = self._ready() + with pytest.raises(ValueError, match="not supported"): + await c.run_code("console.log('x')", language="javascript") + await c.close() + + +# ───────────────────────── Async _ensure_kernel parity ─────────────────────── + + +@respx.mock +def _make_async_capsule(capsule_id: str = "sb-1") -> AsyncCapsule: + """Construct an AsyncCapsule without going through ``create()``.""" + from wrenn.client import AsyncWrennClient + from wrenn.models import Capsule as CapsuleModel + + client = AsyncWrennClient(api_key=API_KEY, base_url=BASE) + info = CapsuleModel(id=capsule_id) + return AsyncCapsule(_capsule_id=capsule_id, _client=client, _info=info) + + +class TestAsyncEnsureKernel: + @pytest.mark.asyncio + @respx.mock + async def test_async_creates_kernel_when_none_exist(self): + c = _make_async_capsule() + proxy_base = "https://8888-sb-1.app.wrenn.dev" + list_route = respx.get(f"{proxy_base}/api/kernels").respond(200, json=[]) + create_route = respx.post(f"{proxy_base}/api/kernels").respond( + 201, json={"id": "k-new", "name": "wrenn"} + ) + kid = await c._ensure_kernel() + assert kid == "k-new" + body = json.loads(create_route.calls[0].request.content) + assert body == {"name": "wrenn"} + assert list_route.called + await c.close() + + @pytest.mark.asyncio + @respx.mock + async def test_async_reuses_existing_wrenn_kernel(self): + c = _make_async_capsule() + proxy_base = "https://8888-sb-1.app.wrenn.dev" + respx.get(f"{proxy_base}/api/kernels").respond( + 200, + json=[ + {"id": "k-other", "name": "python3"}, + {"id": "k-wrenn", "name": "wrenn"}, + ], + ) + create = respx.post(f"{proxy_base}/api/kernels").respond(201, json={}) + kid = await c._ensure_kernel() + assert kid == "k-wrenn" + assert not create.called + await c.close() + + @pytest.mark.asyncio + @respx.mock + async def test_async_retries_on_5xx_then_succeeds(self): + c = _make_async_capsule() + proxy_base = "https://8888-sb-1.app.wrenn.dev" + responses = [ + httpx.Response(503), + httpx.Response(200, json=[{"id": "k-1", "name": "wrenn"}]), + ] + respx.get(f"{proxy_base}/api/kernels").mock(side_effect=responses) + with patch("asyncio.sleep") as sleep_mock: + + async def _noop(_s): + return None + + sleep_mock.side_effect = _noop + kid = await c._ensure_kernel(jupyter_timeout=5) + assert kid == "k-1" + await c.close() + + @pytest.mark.asyncio + @respx.mock + async def test_async_raises_on_4xx(self): + c = _make_async_capsule() + proxy_base = "https://8888-sb-1.app.wrenn.dev" + respx.get(f"{proxy_base}/api/kernels").respond(401) + with pytest.raises(httpx.HTTPStatusError): + await c._ensure_kernel(jupyter_timeout=2) + await c.close() + + @pytest.mark.asyncio + @respx.mock + async def test_async_caches_kernel_id(self): + c = _make_async_capsule() + proxy_base = "https://8888-sb-1.app.wrenn.dev" + route = respx.get(f"{proxy_base}/api/kernels").respond( + 200, json=[{"id": "k-1", "name": "wrenn"}] + ) + await c._ensure_kernel() + await c._ensure_kernel() + assert route.call_count == 1 + await c.close()