diff --git a/contree_cli/cli/operation.py b/contree_cli/cli/operation.py index 1d09697..5fbc35b 100644 --- a/contree_cli/cli/operation.py +++ b/contree_cli/cli/operation.py @@ -13,6 +13,7 @@ from __future__ import annotations import argparse +import contextlib import itertools import json import logging @@ -21,7 +22,7 @@ from datetime import datetime from typing import Any -from contree_cli import CLIENT, FORMATTER, ArgumentsProtocol, SetupResult +from contree_cli import CLIENT, FORMATTER, SESSION_STORE, ArgumentsProtocol, SetupResult from contree_cli.cli.show import ShowArgs, cmd_show from contree_cli.client import ApiError, ContreeClient, PaginatedFetcher from contree_cli.output import OutputFormatter @@ -496,11 +497,23 @@ def cmd_wait(args: WaitArgs) -> int | None: exit_status = 0 sleep_time = 0.5 + # Op cache is best-effort — `wait` doesn't strictly need the + # session store, and the cmd is reachable from contexts that don't + # set one (e.g. standalone scripts). + try: + store = SESSION_STORE.get() + except LookupError: + store = None + while pending and time.monotonic() < deadline: for uuid in list(pending): resp = client.get(f"/v1/operations/{uuid}") op = json.loads(resp.read()) + if op.get("status") in TERMINAL_STATUSES: + if store is not None: + with contextlib.suppress(Exception): + store.cache[(uuid, "operation")] = op pending.discard(uuid) exit_code = extract_exit_code(op) formatter( diff --git a/contree_cli/cli/run.py b/contree_cli/cli/run.py index 74e74d4..ab07273 100644 --- a/contree_cli/cli/run.py +++ b/contree_cli/cli/run.py @@ -48,6 +48,7 @@ import argparse import base64 +import contextlib import fnmatch import functools import io @@ -62,9 +63,18 @@ import uuid from dataclasses import dataclass, field from multiprocessing.pool import ThreadPool +from typing import Any from contree_cli import CLIENT, FORMATTER, SESSION_STORE, ArgumentsProtocol, SetupResult -from contree_cli.client import ApiError, ContreeClient, decode_stream, resolve_image +from contree_cli.client import ( + RETRYABLE_NETWORK_ERRORS, + ApiError, + ContreeClient, + decode_event_chunk, + decode_stream, + iter_sse_events, + resolve_image, +) from contree_cli.mapped_file import MAPPING_RULES, MappedFile from contree_cli.output import ( DefaultFormatter, @@ -505,11 +515,227 @@ def _build_payload( return payload +TERMINAL_OP_STATUSES = frozenset({"SUCCESS", "FAILED", "CANCELLED"}) + + +@dataclass +class TerminalSummary: + """Authoritative end-of-stream snapshot built from the SSE events + themselves — `cmd_run` consumes this directly instead of doing a + second ``GET /operations/{uuid}``. + + If SSE ends without a `completion` frame but a between-attempt GET + detects the op is already terminal, the full op dict lands in + ``fallback_op`` so `cmd_run` can use it directly.""" + + completion: dict[str, Any] | None = None + exit_event: dict[str, Any] | None = None + stdout: bytearray = field(default_factory=bytearray) + stderr: bytearray = field(default_factory=bytearray) + fallback_op: dict[str, Any] | None = None + + +def _check_terminal_via_get( + client: ContreeClient, op_uuid: str, summary: TerminalSummary +) -> bool: + """Poll `GET /v1/operations/{uuid}`; if the op is in a terminal + status, park the full response on ``summary.fallback_op`` and + return True. Any GET failure or non-terminal status returns False + so the caller keeps retrying SSE.""" + try: + resp = client.request("GET", f"/v1/operations/{op_uuid}") + op = json.loads(resp.read()) + except (ApiError, ValueError) as exc: + logger.debug("terminal check GET failed: %s", exc) + return False + except RETRYABLE_NETWORK_ERRORS as exc: + logger.debug("terminal check GET failed: %s", exc) + return False + if isinstance(op, dict) and op.get("status") in TERMINAL_OP_STATUSES: + summary.fallback_op = op + return True + return False + + +TIGHT_LOOP_FLOOR = 0.5 + + +def _stream_events_until_close( + client: ContreeClient, + op_uuid: str, + formatter: OutputFormatter, +) -> TerminalSummary: + """Open `follow=1` SSE for *op_uuid* and write events to stdio, + transparently resuming on network drops / mid-stream errors using + ``Last-Event-Id`` for replay-free continuation. + + For ``DefaultFormatter``: ``stdout`` / ``stderr`` events go to + ``sys.std*.buffer`` directly so the user sees output as it arrives. + For JSON formatters: data goes to ``log.debug`` so the JSON output + isn't polluted, but the chunks are accumulated into the returned + ``TerminalSummary`` so the caller can render full stdout/stderr + without a follow-up GET. + + Loops until one of: `completion` event received, GET-detected + terminal status, ``BrokenPipeError`` from a local stdio write, or + ``KeyboardInterrupt``. Backoff between attempts is delegated to + ``client.request`` — every SSE reconnect goes through its + ``RETRY_DELAYS`` ladder before raising, so an extra streamer-level + ramp would double up. The one floor kept here (``TIGHT_LOOP_FLOOR``) + only kicks in when a cycle made no forward progress, guarding + against a server that returns immediate empty streams for an + executing op. + + ``BrokenPipeError`` from a local stdio write propagates unchanged — + it means the shell pipe closed and retrying cannot help; the + caller cancels the op and exits. + """ + is_default = isinstance(formatter, DefaultFormatter) + last_id: int = -1 + summary = TerminalSummary() + + while True: + headers: dict[str, str] | None = None + if last_id >= 0: + headers = {"Last-Event-Id": str(last_id)} + try: + resp = client.request( + "GET", + f"/v1/operations/{op_uuid}/events?follow=1", + headers=headers, + ) + except ApiError as exc: + logger.debug("SSE connect failed after client retries: %s", exc) + if _check_terminal_via_get(client, op_uuid, summary): + return summary + continue + except RETRYABLE_NETWORK_ERRORS as exc: + logger.debug("SSE connect network error after client retries: %s", exc) + if _check_terminal_via_get(client, op_uuid, summary): + return summary + continue + + events_before = last_id + try: + for ev in iter_sse_events(resp): + ev_id = ev.get("id") + if isinstance(ev_id, int): + last_id = ev_id + ev_type = ev.get("type") + data = ev.get("data") + match ev_type: + case "stdout": + chunk = decode_event_chunk(data) + summary.stdout.extend(chunk) + if is_default: + sys.stdout.buffer.write(chunk) + sys.stdout.buffer.flush() + case "stderr": + chunk = decode_event_chunk(data) + summary.stderr.extend(chunk) + if is_default: + sys.stderr.buffer.write(chunk) + sys.stderr.buffer.flush() + case "exit": + # spid=1 is the main process — its exit code/timed_out + # drive the CLI's own exit code. + if ev.get("spid") == 1: + summary.exit_event = ev + logger.debug("event: %s", ev) + case "sse_error": + logger.warning( + "server-side stream error (last_id=%s): %s", + last_id, + ev.get("message"), + ) + case "completion": + # Authoritative terminal frame — don't wait for the server + # to close, return the summary for the caller. + summary.completion = ev + logger.debug("event: %s", ev) + return summary + case _: + logger.debug("event: %s", ev) + except BrokenPipeError: + # Local stdout/stderr was closed by the shell (e.g. piping + # into `head`). Retrying cannot help — the caller cancels + # the op and exits. + raise + except RETRYABLE_NETWORK_ERRORS as exc: + logger.debug("SSE stream broken (last_id=%s): %s", last_id, exc) + finally: + with contextlib.suppress(Exception): + resp.close() + + if _check_terminal_via_get(client, op_uuid, summary): + return summary + + # No forward progress this cycle → briefly floor the loop so a + # server that keeps returning immediate EOFs doesn't spin us. + if last_id == events_before: + time.sleep(TIGHT_LOOP_FLOOR) + + +def _build_op_from_summary(op_uuid: str, summary: TerminalSummary) -> dict[str, Any]: + """Synthesize a `GET /operations/{uuid}` shape from the SSE terminal + summary — completion event drives status / error / duration / image + metadata; exit event drives exit_code + timed_out; stdout/stderr + are reassembled from the streamed chunks. + + Lets `cmd_run` avoid a second HTTP round-trip on the happy path + while keeping the downstream `_display_operation` consumers + (default and JSON formatters) untouched.""" + assert summary.completion is not None + completion_data = summary.completion.get("data") or {} + status = completion_data.get("status") + state: dict[str, Any] = {} + if summary.exit_event: + exit_data = summary.exit_event.get("data") or {} + if "code" in exit_data: + state["exit_code"] = int(exit_data["code"]) + if "timed_out" in exit_data: + state["timed_out"] = bool(exit_data["timed_out"]) + result_image_uuid = completion_data.get("result_image_uuid") + return { + "uuid": op_uuid, + "kind": "instance", + "status": status, + "error": completion_data.get("error"), + "duration": (completion_data.get("duration_ms") or 0) / 1000.0, + "image_size": completion_data.get("image_size_bytes"), + "result_image_uuid": result_image_uuid, + "metadata": { + "result": { + "stdout": { + "value": bytes(summary.stdout).decode("utf-8", errors="replace"), + "encoding": "ascii", + "truncated": False, + }, + "stderr": { + "value": bytes(summary.stderr).decode("utf-8", errors="replace"), + "encoding": "ascii", + "truncated": False, + }, + "state": state or None, + } + }, + "result": {"image": result_image_uuid, "tag": None}, + } + + def _display_operation( op: dict[str, object], formatter: OutputFormatter, + live_streamed: bool = False, ) -> None: - """Display an operation result using the given formatter.""" + """Display an operation result using the given formatter. + + ``live_streamed=True`` means the SSE path already wrote stdout/ + stderr to stdio incrementally — DefaultFormatter then has nothing + left to print. ``live_streamed=False`` is the GET-fallback path + (legacy server or aborted stream) — DefaultFormatter prints + sanitized stdout/stderr from the operation's metadata blob. + """ result = op.get("result") or {} assert isinstance(result, dict) metadata = op.get("metadata") or {} @@ -544,10 +770,14 @@ def _display_operation( "only json/json-pretty/default are supported" ) - # For DefaultFormatter, just only print stdout/stderr + if live_streamed: + # stdout/stderr already streamed live via SSE — nothing more to print. + return + + # Legacy / fallback path (no completion event received): print + # sanitized stdout/stderr from op metadata. stdout = _BREAKING_ESC_RE.sub("", decode_stream(instance_result.get("stdout"))) stderr = _BREAKING_ESC_RE.sub("", decode_stream(instance_result.get("stderr"))) - if stdout: sys.stdout.write(stdout) if not stdout.endswith("\n"): @@ -688,19 +918,27 @@ def _norm(item: object) -> dict[str, object]: formatter.flush() return None - # 5. Poll until terminal status - sleep_time = 0.5 + # 5. Stream events (follow=1) — write stdout/stderr to stdtio as they + # come, log other events at debug, accumulate the terminal frame. + store = SESSION_STORE.get() try: - time.sleep(sleep_time) - sleep_time += sleep_time - while True: + summary = _stream_events_until_close(client, op_uuid, formatter) + if summary.completion is not None: + # Authoritative terminal frame from the server — build the + # full op dict from the SSE events themselves, no GET. + op = _build_op_from_summary(op_uuid, summary) + elif summary.fallback_op is not None: + # SSE couldn't deliver `completion`, but a GET between + # retries confirmed the op is terminal — use that dict. + op = summary.fallback_op + else: + # Safety net: streamer normally loops until either + # completion or a terminal GET, so this path is only hit + # in tests where the stub queue drains early. resp = client.get(f"/v1/operations/{op_uuid}") op = json.loads(resp.read()) - if op["status"] in TERMINAL_STATUSES: - break - time.sleep(sleep_time) - if sleep_time < 5: - sleep_time += sleep_time + cache_key = (op_uuid, "operation") + store.cache[cache_key] = op except KeyboardInterrupt: try: client.delete(f"/v1/operations/{op_uuid}") @@ -708,6 +946,17 @@ def _norm(item: object) -> dict[str, object]: except (ApiError, KeyboardInterrupt, OSError): pass raise + except BrokenPipeError: + # Local stdout/stderr was closed (e.g. `contree run | head`). + # Cancel the op, silence further stdio writes, then exit 141 + # so callers see the SIGPIPE convention (128 + 13). + with contextlib.suppress(ApiError, OSError): + client.delete(f"/v1/operations/{op_uuid}") + with contextlib.suppress(OSError): + devnull = os.open(os.devnull, os.O_WRONLY) + os.dup2(devnull, sys.stdout.fileno()) + os.close(devnull) + raise SystemExit(141) from None # 6. Cache terminal operation result store.cache[(op_uuid, "operation")] = op @@ -735,7 +984,7 @@ def _norm(item: object) -> dict[str, object]: ) # 7. Display result - _display_operation(op, formatter) + _display_operation(op, formatter, live_streamed=summary.completion is not None) result = op.get("result") or {} assert isinstance(result, dict) diff --git a/contree_cli/cli/show.py b/contree_cli/cli/show.py index 4b59aa2..8944f5d 100644 --- a/contree_cli/cli/show.py +++ b/contree_cli/cli/show.py @@ -61,7 +61,7 @@ def cmd_show(args: ShowArgs) -> int | None: cache_key = (op_uuid, "operation") cached = store.cache.get(cache_key) - if cached is not None: + if isinstance(cached, dict) and cached.get("status") in TERMINAL: op = cast(dict[str, Any], cached) else: resp = client.get(f"/v1/operations/{op_uuid}") diff --git a/contree_cli/client.py b/contree_cli/client.py index 6648acc..5bf15dc 100644 --- a/contree_cli/client.py +++ b/contree_cli/client.py @@ -1,6 +1,7 @@ from __future__ import annotations import base64 +import binascii import collections import contextlib import http.client @@ -12,10 +13,12 @@ import sys import threading import time +import zlib from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Iterator from concurrent.futures import Future, ThreadPoolExecutor -from importlib.metadata import PackageNotFoundError, version +from contextlib import suppress +from importlib.metadata import PackageNotFoundError, distribution from typing import IO, Any, cast from urllib.parse import urlencode, urlsplit @@ -23,7 +26,33 @@ log = logging.getLogger(__name__) -RETRY_DELAYS = (1, 2, 4, 5, 10, 10, 10) +RETRY_DELAYS = (0.1, 0.2, 0.5, 1, 2, 5) + + +def retry_generator() -> Iterator[float]: + """Return a primed generator that yields successive backoff delays. + + The first ``next(g)`` returns ``RETRY_DELAYS[0]``; subsequent calls + walk the ladder and, once exhausted, keep yielding the tail delay + forever. Callers don't have to guard against ``StopIteration`` — + ``next(g)`` always returns a valid sleep time — and don't have to + do any index arithmetic to derive it. A caller that wants a + finite retry budget bounds it externally (e.g. an attempt counter).""" + + def gen() -> Iterator[float]: + # `yield 0` gives a caller who *doesn't* prime a clean way to + # sleep uniformly before every attempt (including the first, + # which is a no-op). We prime here so callers who don't want + # that get the real first delay straight away. + yield 0 + yield from RETRY_DELAYS + while True: + yield RETRY_DELAYS[-1] + + g = gen() + next(g) + return g + # Socket-level / connection-level errors that warrant a retry. DNS hiccups # (gaierror), refused/reset connections, and broken HTTP framing are all @@ -38,9 +67,17 @@ def cli_version() -> str: try: - return version("contree-cli") + dist = distribution("contree-cli") except PackageNotFoundError: return "editable" + raw = dist.read_text("direct_url.json") + if raw: + try: + if json.loads(raw).get("dir_info", {}).get("editable"): + return "editable" + except ValueError: + pass + return dist.version CLI_USER_AGENT = ( @@ -138,6 +175,88 @@ def format_bytes(self, data: bytes) -> str: return f"" +class GzipResponse: + """Wrap an HTTPResponse, inflating `Content-Encoding: gzip` via + `zlib.decompressobj(wbits=31)` so Z_SYNC_FLUSH'd SSE frames decode + incrementally without buffering whole responses.""" + + def __init__(self, resp: http.client.HTTPResponse) -> None: + self.resp = resp + self.decomp = zlib.decompressobj(wbits=31) + self.buf = bytearray() + self.eof = False + + def pump(self, want: int = -1) -> None: + while not self.eof and (want < 0 or len(self.buf) < want): + # read1() returns whatever's in the socket buffer after one + # underlying read; plain read(n) blocks for n bytes, which + # stalls SSE streaming because tiny frames never fill it. + chunk = self.resp.read1(8192) + if not chunk: + tail = self.decomp.flush() + if tail: + self.buf.extend(tail) + self.eof = True + return + decoded = self.decomp.decompress(chunk) + if decoded: + self.buf.extend(decoded) + + def read(self, amt: int | None = None) -> bytes: + if amt is None or amt < 0: + self.pump() + out = bytes(self.buf) + self.buf.clear() + return out + self.pump(amt) + out = bytes(self.buf[:amt]) + del self.buf[:amt] + return out + + def readline(self, limit: int = -1) -> bytes: + while True: + nl = self.buf.find(b"\n") + if nl >= 0: + end = nl + 1 + if 0 <= limit < end: + end = limit + out = bytes(self.buf[:end]) + del self.buf[:end] + return out + if self.eof: + out = bytes(self.buf) + self.buf.clear() + return out + self.pump(want=len(self.buf) + 1) + + def close(self) -> None: + with contextlib.suppress(Exception): + self.resp.close() + + def getheader(self, name: str, default: str | None = None) -> str | None: + # Hide Content-Encoding so downstream readers don't try a + # second decompression; Content-Length is bogus post-inflate. + lname = name.lower() + if lname in ("content-encoding", "content-length"): + return default + return self.resp.getheader(name, default) + + def getheaders(self) -> list[tuple[str, str]]: + return [ + (k, v) + for k, v in self.resp.getheaders() + if k.lower() not in ("content-encoding", "content-length") + ] + + @property + def status(self) -> int: + return self.resp.status + + @property + def reason(self) -> str: + return self.resp.reason + + class BufferedResponse: """Replay an HTTPResponse from buffered bytes (for debug body logging).""" @@ -221,6 +340,7 @@ def request( merged: dict[str, str] = { **self._build_headers(), "User-Agent": CLI_USER_AGENT, + "Accept-Encoding": "gzip", } if body is not None: merged.setdefault("Content-Type", "application/json") @@ -228,9 +348,6 @@ def request( merged.update(headers) full_path = self._prefix + path - last_error: ApiError | None = None - last_network_error: BaseException | None = None - attempts = len(RETRY_DELAYS) + 1 log.debug( "%s %s headers=%s body=%s", @@ -240,25 +357,21 @@ def request( BodyFormatter(body, content_type=merged.get("Content-Type", "")), ) - for attempt in range(attempts): - if last_error is not None or last_network_error is not None: - delay = RETRY_DELAYS[attempt - 1] - if last_network_error is not None: - log.warning( - "Network error (%s), retrying in %ds…", - type(last_network_error).__name__, - delay, - ) - else: - assert last_error is not None - log.warning( - "Server error %d, retrying in %ds…", - last_error.status, - delay, - ) - time.sleep(delay) - - if attempt > 0 and hasattr(body, "seek"): + # Backoff delays come from a primed `retry_generator()` — each + # failure just calls `next(delays)` and gets the next value. + # The generator yields the tail delay forever, and there is no + # outer retry budget: transient errors are retried until either + # the server returns a non-retryable response or the user hits + # Ctrl+C. Non-retryable 4xx responses raise immediately. + delays = retry_generator() + + def sleep_for_retry(reason_fmt: str, *reason_args: object) -> None: + """Pull the next backoff delay, log the reason, sleep, and + rewind the request body if it's seekable.""" + delay = next(delays) + log.warning(reason_fmt + " retrying in %ss…", *reason_args, delay) + time.sleep(delay) + if body is not None and hasattr(body, "seek"): stream = cast(IO[bytes], body) if not stream.seekable(): raise ApiError( @@ -267,6 +380,8 @@ def request( "Cannot retry: streaming body is not seekable", ) stream.seek(0) + + while True: try: conn = self._connect() conn.request(method, full_path, body, merged) @@ -278,13 +393,11 @@ def request( # would just spin through the back-off ladder for nothing. raise except RETRYABLE_NETWORK_ERRORS as exc: - last_network_error = exc - last_error = None + sleep_for_retry("Network error (%s),", type(exc).__name__) continue - # Successful round-trip clears the network-error trail so the - # final raise below doesn't pick up stale failure context. - last_network_error = None + if (resp.getheader("Content-Encoding", "") or "").lower() == "gzip": + resp = cast(http.client.HTTPResponse, GzipResponse(resp)) if 200 <= resp.status < 300: log.debug( @@ -321,17 +434,12 @@ def request( ) error = ApiError(resp.status, resp.reason, resp_body) - if 500 <= resp.status < 600: - last_error = error + if resp.status in (410, 425) or 500 <= resp.status < 600: + sleep_for_retry("Server error %d,", resp.status) continue raise error - if last_network_error is not None: - raise last_network_error - assert last_error is not None - raise last_error - def log_and_buffer( self, method: str, @@ -340,8 +448,10 @@ def log_and_buffer( ) -> http.client.HTTPResponse: """Read & log a textual response body; pass binary streams through.""" content_type = resp.getheader("Content-Type", "") or "" + # event-stream is line-streamed and unbounded — buffering it would block. + is_event_stream = "event-stream" in content_type textual = not content_type or "json" in content_type or "text" in content_type - if not textual: + if is_event_stream or not textual: log.debug( "%s %s response body: ", method, @@ -535,6 +645,73 @@ def stream_response( yield chunk +def iter_sse_events(resp: http.client.HTTPResponse) -> Iterator[dict[str, object]]: + """Yield one dict per SSE frame. + + Normal frames carry a JSON object in ``data:`` — that object is + yielded as-is. Server-pushed error frames use ``event: sse_error`` + with a plain-text ``data:`` body (the exception message) — those + surface as ``{"type": "sse_error", "message": , "id": ...}`` + so callers can log + decide whether to reconnect with + ``Last-Event-Id`` set to the last good id. + """ + data_lines: list[str] = [] + event_name: str | None = None + event_id: str | None = None + + def emit() -> Iterator[dict[str, object]]: + nonlocal event_name, event_id + if not data_lines: + return + body = "\n".join(data_lines) + if event_name == "sse_error": + payload: dict[str, object] = {"type": "sse_error", "message": body} + if event_id is not None: + payload["id"] = event_id + yield payload + else: + with suppress(json.JSONDecodeError): + decoded = json.loads(body) + if isinstance(decoded, dict): + yield decoded + data_lines.clear() + event_name = None + event_id = None + + while True: + line_bytes = resp.readline() + if not line_bytes: + yield from emit() + return + line = line_bytes.decode("utf-8", errors="replace").rstrip("\r\n") + match line: + case "": + yield from emit() + case _ if line.startswith(":"): + pass # SSE comment / keepalive + case _ if line.startswith("data:"): + data_lines.append(line[5:].lstrip(" ")) + case _ if line.startswith("event:"): + event_name = line[6:].lstrip(" ").strip() or None + case _ if line.startswith("id:"): + event_id = line[3:].lstrip(" ").strip() or None + + +def decode_event_chunk(data: object) -> bytes: + """Decode `{value, encoding}` event payload to raw bytes (ascii or base64).""" + if not isinstance(data, dict): + return b"" + value = data.get("value", "") + if not isinstance(value, str) or not value: + return b"" + encoding = data.get("encoding", "ascii") + if encoding == "base64": + with suppress(binascii.Error, ValueError): + return base64.b64decode(value) + return b"" + return value.encode("utf-8", errors="replace") + + def decode_stream(stream: dict[str, object] | None) -> str: """Decode an API StreamRepr object to a string.""" if not stream: diff --git a/contree_cli/docker/kw_run.py b/contree_cli/docker/kw_run.py index 8bc72fe..53ee6f1 100644 --- a/contree_cli/docker/kw_run.py +++ b/contree_cli/docker/kw_run.py @@ -5,19 +5,22 @@ import json import logging import shlex -import time from dataclasses import dataclass, field -from typing import ClassVar +from typing import Any, ClassVar +from contree_cli.cli.run import ( + TERMINAL_OP_STATUSES, + _build_op_from_summary, + _stream_events_until_close, +) from contree_cli.client import decode_stream +from contree_cli.output import DefaultFormatter from .context import BuildContext from .keyword import DockerKeyword, parse_command_form logger = logging.getLogger(__name__) -TERMINAL_STATUSES = frozenset({"SUCCESS", "FAILED", "CANCELLED"}) - @dataclass(frozen=True, repr=False) class RunKeyword(DockerKeyword): @@ -91,33 +94,42 @@ def spawn(self, ctx: BuildContext, parts: tuple[str, ...]) -> tuple[str, str]: payload["files"] = ctx.pending_files_payload() resp = ctx.client.post_json("/v1/instances", payload) - op = json.loads(resp.read()) - op_uuid: str = op["uuid"] + spawn_op = json.loads(resp.read()) + op_uuid: str = spawn_op["uuid"] logger.info( "RUN spawned op=%s: %s", op_uuid, display_title(parts, self.shell_form) ) - op = poll(ctx, op_uuid) + op = stream_and_resolve(ctx, op_uuid) check_success(op, parts, self.shell_form) result = op.get("result") or {} assert isinstance(result, dict) new_image = result.get("image") if not new_image: raise RuntimeError("RUN succeeded but no image was produced") - log_streams(op) return str(new_image), op_uuid -def poll(ctx: BuildContext, op_uuid: str) -> dict[str, object]: - delay = 0.5 - while True: - time.sleep(delay) - resp = ctx.client.get(f"/v1/operations/{op_uuid}") - op = json.loads(resp.read()) - if op["status"] in TERMINAL_STATUSES: - return op # type: ignore[no-any-return] - if delay < 5: - delay += delay +def stream_and_resolve(ctx: BuildContext, op_uuid: str) -> dict[str, Any]: + """Stream RUN output live (docker-style), then materialise the + full op dict. + + Preference order for the op payload matches ``cmd_run``: + ``completion`` event → SSE `fallback_op` from a terminal GET → + safety-net GET. When we fell back to the plain endpoint (no + live stream), replay the captured stdout/stderr from the op's + metadata so build users still see what the RUN produced.""" + summary = _stream_events_until_close(ctx.client, op_uuid, DefaultFormatter()) + if summary.completion is not None: + return _build_op_from_summary(op_uuid, summary) + if summary.fallback_op is not None: + log_streams(summary.fallback_op) + return summary.fallback_op + resp = ctx.client.get(f"/v1/operations/{op_uuid}") + op: dict[str, Any] = json.loads(resp.read()) + if op.get("status") in TERMINAL_OP_STATUSES: + log_streams(op) + return op def check_success( diff --git a/contree_cli/docker/url_fetch.py b/contree_cli/docker/url_fetch.py index 5dda18c..84e1c35 100644 --- a/contree_cli/docker/url_fetch.py +++ b/contree_cli/docker/url_fetch.py @@ -23,7 +23,6 @@ import urllib.parse import urllib.request from dataclasses import dataclass -from typing import IO, cast from contree_cli.client import ApiError, ContreeClient from contree_cli.session import SessionStore @@ -79,27 +78,25 @@ def fetch_and_upload( assert source is not None try: - reader = HashingReader(source) - upload_headers: dict[str, str] = {"Content-Type": "application/octet-stream"} - content_length = parse_content_length(response_headers) - if content_length > 0: - upload_headers["Content-Length"] = str(content_length) - - resp = client.request( - "POST", - "/v1/files", - body=cast(IO[bytes], reader), - headers=upload_headers, - ) - data = json.loads(resp.read()) + body_bytes = source.read() # type: ignore[attr-defined] finally: close = getattr(source, "close", None) if callable(close): close() + assert isinstance(body_bytes, (bytes, bytearray)) + body_bytes = bytes(body_bytes) + + resp = client.request( + "POST", + "/v1/files", + body=body_bytes, + headers={"Content-Type": "application/octet-stream"}, + ) + data = json.loads(resp.read()) file_uuid = str(data["uuid"]) - sha = reader.hasher.hexdigest() - size = reader.bytes_read + sha = hashlib.sha256(body_bytes).hexdigest() + size = len(body_bytes) write_metadata( store, @@ -195,38 +192,6 @@ def conditional_headers(meta: dict[str, object]) -> dict[str, str]: return headers -def parse_content_length(headers: dict[str, str]) -> int: - raw = headers.get("content-length") - if not raw: - return -1 - try: - return int(raw) - except ValueError: - return -1 - - -class HashingReader: - """Read-only adapter that hashes bytes as they flow through. - - Intentionally has no ``seek`` attribute so ``ContreeClient.request`` - does not attempt to rewind the upstream HTTP body on retry. - """ - - __slots__ = ("bytes_read", "hasher", "source") - - def __init__(self, source: object) -> None: - self.source = source - self.hasher = hashlib.sha256() - self.bytes_read = 0 - - def read(self, amt: int | None = None) -> bytes: - chunk = self.source.read() if amt is None else self.source.read(amt) # type: ignore[attr-defined] - if chunk: - self.hasher.update(chunk) - self.bytes_read += len(chunk) - return chunk # type: ignore[no-any-return] - - def http_head(url: str, *, timeout: int = DOWNLOAD_TIMEOUT_DEFAULT) -> dict[str, str]: """Best-effort ``HEAD`` request. Returns empty headers on any failure. diff --git a/tests/conftest.py b/tests/conftest.py index 34828ca..faf4332 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -77,6 +77,16 @@ def __post_init__(self) -> None: def read(self, amt: int | None = None) -> bytes: return self.body + def readline(self, size: int = -1) -> bytes: + """Empty bytes signals EOF to `iter_sse_events`, which makes any + SSE attempt against a non-SSE mock no-op out without raising — + the test then exercises the GET fallback as before.""" + return b"" + + def read1(self, amt: int | None = None) -> bytes: + """Match BufferedReader.read1 for streaming-style consumers.""" + return self.body + def getheader(self, name: str, default: str | None = None) -> str | None: assert self.headers is not None for key, value in self.headers.items(): @@ -113,6 +123,7 @@ class FakeConnection: def __init__(self) -> None: self.requests: list[RecordedRequest] = [] self.responses: list[FakeResponse] = [] + self._last_path: str = "" def request( self, @@ -131,8 +142,20 @@ def request( body = b"".join(chunks) recorded = body if isinstance(body, (bytes, bytearray, type(None))) else None self.requests.append(RecordedRequest(method, path, recorded, headers or {})) + self._last_path = path def getresponse(self) -> FakeResponse: + # `/events?follow=1` is the modern wait path — auto-serve an empty + # SSE response so the GET-/operations mock the test queued up + # isn't consumed by the events probe. Real tests that want to + # exercise the SSE wire format must mock the connection + # differently anyway. + if "/events" in self._last_path: + return FakeResponse( + status=200, + body=b"", + headers={"Content-Type": "text/event-stream"}, + ) return self.responses.pop(0) diff --git a/tests/test_build.py b/tests/test_build.py index fb54b8c..67f3f84 100644 --- a/tests/test_build.py +++ b/tests/test_build.py @@ -65,7 +65,7 @@ def run_build( SESSION_STORE.set(SessionStore(db_path, "placeholder")) ctx = copy_context() with ( - patch("contree_cli.docker.kw_run.time.sleep"), + patch("contree_cli.cli.run.time.sleep"), patch("contree_cli.docker.kw_from.time.sleep"), ): return ctx.run(cmd_build, args) @@ -121,7 +121,7 @@ def test_build_arg_namespace_decodes_to_build_args(self): class TestSimpleBuild: - def test_from_run_creates_two_api_calls(self, context_dir, db_path): + def test_from_run_creates_expected_api_calls(self, context_dir, db_path): write_dockerfile( context_dir, "FROM tag:ubuntu:latest\nRUN echo hi\n", @@ -135,13 +135,20 @@ def test_from_run_creates_two_api_calls(self, context_dir, db_path): ] rc = run_build(tc, args, responses, db_path) assert rc is None - assert tc.request_count == 3 + # 4 wire calls: FROM's tag lookup, POST instances, SSE events + # follow, then the streamer's terminal GET. + assert tc.request_count == 4 assert tc.get_request(0).method == "GET" assert "/v1/images" in tc.get_request(0).path assert tc.get_request(1).method == "POST" assert "/v1/instances" in tc.get_request(1).path assert tc.get_request(2).method == "GET" - assert "/v1/operations" in tc.get_request(2).path + assert "/v1/operations/op-1/events?follow=1" in tc.get_request(2).path + assert tc.get_request(3).method == "GET" + assert ( + tc.get_request(3).path[-len("/v1/operations/op-1") :] + == "/v1/operations/op-1" + ) def test_run_payload_carries_command(self, context_dir, db_path): write_dockerfile( @@ -166,6 +173,52 @@ def test_run_payload_carries_command(self, context_dir, db_path): assert body["command"] == "apt-get update" assert body["shell"] is True + def test_run_streams_stdout_live(self, context_dir, db_path, capsys): + """Docker-compat mode: RUN output must reach the user's terminal + as the SSE events arrive, not just after the op completes. + + Uses a stubbed streamer that writes a chunk to `sys.stdout.buffer` + and returns a completion-populated summary so the build finishes + with the streamed image — mirrors what a live SSE `stdout` frame + followed by a `completion` frame would produce.""" + import sys + + from contree_cli.cli.run import TerminalSummary + + write_dockerfile( + context_dir, + "FROM tag:ubuntu:latest\nRUN echo hi\n", + ) + tc = ContreeTestClient() + + def fake_stream(_client, op_uuid, _formatter): + sys.stdout.buffer.write(b"hi from RUN\n") + sys.stdout.buffer.flush() + summary = TerminalSummary() + summary.completion = { + "type": "completion", + "data": {"status": "SUCCESS", "result_image_uuid": NEW_IMG}, + } + summary.exit_event = {"type": "exit", "spid": 1, "data": {"code": 0}} + summary.stdout.extend(b"hi from RUN\n") + return summary + + with patch( + "contree_cli.docker.kw_run._stream_events_until_close", + side_effect=fake_stream, + ): + rc = run_build( + tc, + BuildArgs(context=str(context_dir)), + [make_tag_lookup(BASE_IMG), make_spawn()], + db_path, + ) + assert rc is None + # The live-streamed chunk lands on stdout before the final + # formatter record — verify both are present. + out = capsys.readouterr().out + assert "hi from RUN" in out + class TestCache: def test_second_build_is_full_cache_hit(self, context_dir, db_path): @@ -227,7 +280,9 @@ def test_no_cache_reruns(self, context_dir, db_path): db_path, ) assert rc is None - assert second.request_count == 3 + # 4 wire calls per build: FROM tag lookup, POST instances, SSE + # follow, and the streamer's terminal GET. + assert second.request_count == 4 def test_no_cache_when_from_layer_is_active_branch(self, context_dir, db_path): """Regression: --no-cache must not blow up when the target layer @@ -431,7 +486,9 @@ def test_final_image_tagged(self, context_dir, db_path): db_path, ) assert rc is None - tag_req = tc.get_request(3) + # PATCH lands after: tag lookup (0), spawn (1), SSE events (2), + # streamer's terminal GET (3). + tag_req = tc.get_request(4) assert tag_req.method == "PATCH" assert NEW_IMG in tag_req.path body = json.loads(tag_req.body.decode()) diff --git a/tests/test_client.py b/tests/test_client.py index 23987a2..47aff44 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,9 +1,12 @@ from __future__ import annotations +import base64 import contextlib +import io import json import logging -from unittest.mock import patch +from importlib.metadata import PackageNotFoundError +from unittest.mock import MagicMock, patch import pytest from conftest import ContreeTestClient, ContreeTestIAMClient, FakeResponse @@ -17,6 +20,9 @@ ContreeJWTClient, HeaderFormatter, PaginatedFetcher, + cli_version, + decode_event_chunk, + iter_sse_events, resolve_image, ) @@ -140,20 +146,25 @@ def test_retries_on_5xx_then_succeeds(self): assert result.body == b'{"ok":true}' mock_sleep.assert_called_once_with(RETRY_DELAYS[0]) - def test_exhausts_retries_then_raises(self): + def test_retries_until_success(self): + """`client.request` retries transient failures without a + budget — only success or a non-retryable status ends the + loop. Verified here by queuing many 500s followed by a 200.""" c = ContreeTestClient("https://contree.dev", "tok") - total = len(RETRY_DELAYS) + 1 - for _ in range(total): + failures = 20 + for _ in range(failures): c.respond(status=500, body=b"err") + c.respond(status=200, body=b'{"ok":true}') - with ( - patch("contree_cli.client.time.sleep") as mock_sleep, - pytest.raises(ApiError) as exc_info, - ): - c.request("GET", "/v1/images") + with patch("contree_cli.client.time.sleep") as mock_sleep: + result = c.request("GET", "/v1/images") - assert exc_info.value.status == 500 - assert mock_sleep.call_count == len(RETRY_DELAYS) + assert result.body == b'{"ok":true}' + assert mock_sleep.call_count == failures + # First failures walk the ladder, the rest reuse the tail delay. + delays = [call.args[0] for call in mock_sleep.call_args_list] + assert delays[: len(RETRY_DELAYS)] == list(RETRY_DELAYS) + assert all(d == RETRY_DELAYS[-1] for d in delays[len(RETRY_DELAYS) :]) def test_no_retry_on_4xx(self): c = ContreeTestClient("https://contree.dev", "tok") @@ -208,22 +219,41 @@ def flaky_connect(): assert call_count["n"] == 3 assert mock_sleep.call_count == 2 - def test_retry_exhausted_raises_network_error(self): - """When retries run out, the last network error propagates.""" - import socket + def test_first_attempt_410_uses_short_delay(self): + """410 on the very first attempt sleeps `RETRY_DELAYS[0]`, not + the wraparound `RETRY_DELAYS[-1]` from the old `attempt-1` index.""" + c = ContreeTestClient("https://contree.dev", "tok") + c.respond(status=410, body=b"gone") + c.respond(status=200, body=b'{"ok":true}') + + with patch("contree_cli.client.time.sleep") as mock_sleep: + c.request("GET", "/v1/images") + delays = [call.args[0] for call in mock_sleep.call_args_list] + assert delays[0] == RETRY_DELAYS[0] + def test_first_attempt_425_uses_short_delay(self): c = ContreeTestClient("https://contree.dev", "tok") + c.respond(status=425, body=b"too early") + c.respond(status=200, body=b'{"ok":true}') - def always_fails(): - raise socket.gaierror(8, "nodename nor servname provided") + with patch("contree_cli.client.time.sleep") as mock_sleep: + c.request("GET", "/v1/images") + delays = [call.args[0] for call in mock_sleep.call_args_list] + assert delays[0] == RETRY_DELAYS[0] - c._connect = always_fails # type: ignore[method-assign] + def test_410_sleeps_once_per_failure(self): + """Regression: 410/425 used to `time.sleep` inside the response + branch AND at the top of the next iteration — doubling the + actual backoff. With the retry-generator refactor each failure + pulls exactly one delay from the ladder.""" + c = ContreeTestClient("https://contree.dev", "tok") + c.respond(status=410, body=b"gone") + c.respond(status=200, body=b'{"ok":true}') - with ( - patch("contree_cli.client.time.sleep"), - pytest.raises(socket.gaierror), - ): + with patch("contree_cli.client.time.sleep") as mock_sleep: c.request("GET", "/v1/images") + assert mock_sleep.call_count == 1 + assert mock_sleep.call_args_list[0].args[0] == RETRY_DELAYS[0] def test_invalid_url_is_not_retried(self): """InvalidURL is a permanent caller-side error — should raise immediately.""" @@ -618,6 +648,62 @@ def test_raises_without_token(self): c.request("GET", "/v1/images") +class TestCliVersion: + """``cli_version()`` must return ``"editable"`` whenever the install is + not a regular wheel — either the package is missing from metadata, or + PEP 610 ``direct_url.json`` marks the install as editable. The update + checker keys off this sentinel to skip PyPI pings during local dev.""" + + def make_dist(self, *, version: str, direct_url: str | None) -> MagicMock: + dist = MagicMock() + dist.version = version + dist.read_text.return_value = direct_url + return dist + + def test_returns_editable_when_package_not_installed(self): + with patch( + "contree_cli.client.distribution", + side_effect=PackageNotFoundError("contree-cli"), + ): + assert cli_version() == "editable" + + def test_returns_editable_for_pep610_editable_install(self): + dist = self.make_dist( + version="0.5.0", + direct_url=json.dumps( + { + "url": "file:///path/to/contree-cli", + "dir_info": {"editable": True}, + }, + ), + ) + with patch("contree_cli.client.distribution", return_value=dist): + assert cli_version() == "editable" + + def test_returns_version_for_regular_install(self): + dist = self.make_dist(version="0.5.0", direct_url=None) + with patch("contree_cli.client.distribution", return_value=dist): + assert cli_version() == "0.5.0" + + def test_returns_version_when_direct_url_lacks_editable_flag(self): + dist = self.make_dist( + version="0.5.0", + direct_url=json.dumps( + { + "url": "https://files.pythonhosted.org/.../contree_cli.whl", + "archive_info": {}, + }, + ), + ) + with patch("contree_cli.client.distribution", return_value=dist): + assert cli_version() == "0.5.0" + + def test_returns_version_when_direct_url_is_malformed(self): + dist = self.make_dist(version="0.5.0", direct_url="not json {") + with patch("contree_cli.client.distribution", return_value=dist): + assert cli_version() == "0.5.0" + + class TestPaginatedFetcherLimit: """``limit=`` lives in :class:`PaginatedFetcher` so callers don't repeat the page-size math, and a small budget like ``--limit 5`` doesn't pull a @@ -679,3 +765,122 @@ def test_small_limit_uses_capped_page_size_in_request(self, contree_client): req = contree_client.get_request(0) assert "limit=6" in req.path assert "limit=1000" not in req.path + + +# --------------------------------------------------------------------------- +# SSE parser: iter_sse_events +# --------------------------------------------------------------------------- + + +def _sse(body: str) -> io.BytesIO: + return io.BytesIO(body.encode("utf-8")) + + +class TestIterSseEvents: + def test_single_frame(self): + body = ( + "id: 1\nevent: stdout\n" + 'data: {"type":"stdout","data":{"value":"hi","encoding":"ascii"}}\n\n' + ) + events = list(iter_sse_events(_sse(body))) + assert events == [ + {"type": "stdout", "data": {"value": "hi", "encoding": "ascii"}} + ] + + def test_multiple_frames(self): + body = ( + 'id: 0\nevent: init\ndata: {"type":"init","data":{}}\n\n' + 'id: 1\nevent: exit\ndata: {"type":"exit","spid":1,"data":{"code":0}}\n\n' + ) + events = list(iter_sse_events(_sse(body))) + assert [e["type"] for e in events] == ["init", "exit"] + assert events[1]["spid"] == 1 + + def test_keepalive_and_blank_lines_ignored(self): + body = ( + ": keepalive\n" + "\n" + ": keepalive again\n" + 'id: 5\nevent: stdout\ndata: {"type":"stdout"}\n\n' + ) + events = list(iter_sse_events(_sse(body))) + assert events == [{"type": "stdout"}] + + def test_sse_error_frame_surfaces_as_dict(self): + body = ( + ": stream ended with error, retry since last event id\n\n" + "event: sse_error\ndata: upstream closed unexpectedly\n\n" + ) + events = list(iter_sse_events(_sse(body))) + assert events == [ + {"type": "sse_error", "message": "upstream closed unexpectedly"} + ] + + def test_sse_error_with_id_carries_id(self): + body = "id: 42\nevent: sse_error\ndata: boom\n\n" + events = list(iter_sse_events(_sse(body))) + assert events == [{"type": "sse_error", "message": "boom", "id": "42"}] + + def test_multiline_data_joined_with_newline(self): + body = 'event: stdout\ndata: {"type":"stdout",\ndata: "value":"x"}\n\n' + events = list(iter_sse_events(_sse(body))) + assert events == [{"type": "stdout", "value": "x"}] + + def test_invalid_json_data_dropped_silently(self): + body = 'event: stdout\ndata: not-json\n\nevent: exit\ndata: {"type":"exit"}\n\n' + events = list(iter_sse_events(_sse(body))) + assert events == [{"type": "exit"}] + + def test_non_dict_json_dropped(self): + body = 'event: stdout\ndata: [1,2,3]\n\nevent: exit\ndata: {"type":"exit"}\n\n' + events = list(iter_sse_events(_sse(body))) + assert events == [{"type": "exit"}] + + def test_frame_at_eof_without_blank_line_still_emits(self): + body = 'event: stdout\ndata: {"type":"stdout"}\n' + events = list(iter_sse_events(_sse(body))) + assert events == [{"type": "stdout"}] + + def test_empty_stream_yields_nothing(self): + events = list(iter_sse_events(_sse(""))) + assert events == [] + + +# --------------------------------------------------------------------------- +# SSE chunk decoder: decode_event_chunk +# --------------------------------------------------------------------------- + + +class TestDecodeEventChunk: + def test_ascii_encoding_default(self): + assert decode_event_chunk({"value": "hello"}) == b"hello" + + def test_explicit_ascii(self): + assert decode_event_chunk({"value": "hi", "encoding": "ascii"}) == b"hi" + + def test_base64_encoding(self): + payload = base64.b64encode(b"\x00\x01\xff").decode("ascii") + assert ( + decode_event_chunk({"value": payload, "encoding": "base64"}) + == b"\x00\x01\xff" + ) + + def test_base64_invalid_returns_empty(self): + assert ( + decode_event_chunk({"value": "!!!not-base64!!!", "encoding": "base64"}) + == b"" + ) + + def test_missing_value_returns_empty(self): + assert decode_event_chunk({}) == b"" + + def test_empty_value_returns_empty(self): + assert decode_event_chunk({"value": ""}) == b"" + + def test_non_dict_returns_empty(self): + assert decode_event_chunk("not a dict") == b"" + assert decode_event_chunk(None) == b"" + assert decode_event_chunk(42) == b"" + + def test_non_string_value_returns_empty(self): + assert decode_event_chunk({"value": 123}) == b"" diff --git a/tests/test_file_cmd.py b/tests/test_file_cmd.py index 2e7af2e..653a885 100644 --- a/tests/test_file_cmd.py +++ b/tests/test_file_cmd.py @@ -236,16 +236,12 @@ def test_non_404_error_propagates( ): session_store.set_image("a1b2c3d4-5678-9abc-def0-111111111111", kind="use") args = FileEditArgs(path="/etc/config.ini") - # Client retries 500s (RETRY_DELAYS has 7 entries -> 8 attempts total) - responses = [ - _api_response({"error": "server error"}, status=500) for _ in range(8) - ] - with ( - patch("time.sleep"), - pytest.raises(ApiError) as exc_info, - ): + # 403 is non-retryable, so the caller sees it on the first hit + # without any retries or sleep. + responses = [_api_response({"error": "forbidden"}, status=403)] + with pytest.raises(ApiError) as exc_info: _run_file_edit(contree_client, args, responses, store=session_store) - assert exc_info.value.status == 500 + assert exc_info.value.status == 403 class TestFileEditNoChanges: diff --git a/tests/test_run.py b/tests/test_run.py index a0b3dbc..04f7a34 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -7,6 +7,7 @@ import os import select from contextvars import copy_context +from typing import ClassVar from unittest.mock import MagicMock, patch import pytest @@ -15,10 +16,13 @@ from contree_cli import CLIENT, FORMATTER, SESSION_STORE from contree_cli.cli.run import ( RunArgs, + TerminalSummary, + _build_op_from_summary, _expand_mapped_files, _is_excluded, _local_file_cache_kind, _read_piped_stdin, + _stream_events_until_close, cmd_run, ) from contree_cli.client import ApiError @@ -96,7 +100,12 @@ def _run_cmd( formatter=None, stdin_mock: MagicMock | None = None, ): - """Run cmd_run with mocked HTTP responses and mocked sleep.""" + """Run cmd_run with mocked HTTP responses and mocked sleep. + + The fake connection auto-serves empty SSE responses for any + GET /events path so existing tests can stay shaped as + `[spawn, op]` without knowing the CLI now opens an SSE first. + """ tc.fake.responses.extend(responses) FORMATTER.set(formatter or JSONFormatter()) @@ -159,11 +168,12 @@ def test_detach_shows_status(self, contree_client, session_store, capsys): class TestPollLoop: def test_poll_until_success(self, contree_client, session_store, capsys): + """SSE follow=1 makes the API serve the terminal state directly — + the CLI no longer polls through intermediate PENDING snapshots.""" session_store.set_image(IMG_UUID, kind="test") args = _default_args() responses = [ _spawn_response(), - _op_response(status="PENDING"), _op_response(status="SUCCESS", exit_code=0), ] rc = _run_cmd(contree_client, args, responses, store=session_store) @@ -321,7 +331,9 @@ class TestCtrlC: def _run_ctrl_c( responses: list[FakeResponse], store: SessionStore ) -> ContreeTestClient: - """Run cmd_run with sleep raising KeyboardInterrupt.""" + """Run cmd_run with the events stream raising KeyboardInterrupt + — simulates the user hitting Ctrl-C while the CLI was waiting + on SSE for the operation to terminate.""" store.set_image(IMG_UUID, kind="test") args = _default_args() tc = ContreeTestClient() @@ -334,7 +346,7 @@ def _run_ctrl_c( with ( patch( - "contree_cli.cli.run.time.sleep", + "contree_cli.cli.run._stream_events_until_close", side_effect=KeyboardInterrupt, ), patch("contree_cli.cli.run.sys.stdin", _tty_stdin()), @@ -344,11 +356,10 @@ def _run_ctrl_c( return tc def test_ctrl_c_cancels_operation(self, session_store): - """On KeyboardInterrupt during sleep, DELETE is sent.""" + """On KeyboardInterrupt during the wait, DELETE is sent.""" tc = self._run_ctrl_c( [ _spawn_response(), - _op_response(status="PENDING"), _api_response({}, status=202), ], session_store, @@ -363,13 +374,71 @@ def test_ctrl_c_delete_failure_still_raises(self, session_store): self._run_ctrl_c( [ _spawn_response(), - _op_response(status="PENDING"), _api_response({"error": "not found"}, status=404), ], session_store, ) +class TestBrokenPipe: + """`BrokenPipeError` from local stdio (shell piped output closed + early, e.g. ``contree run ... | head``) must cancel the remote op + and exit with 141 (128 + SIGPIPE) instead of being misinterpreted + as a remote network drop.""" + + @staticmethod + def _run_broken_pipe( + responses: list[FakeResponse], store: SessionStore + ) -> ContreeTestClient: + store.set_image(IMG_UUID, kind="test") + args = _default_args() + tc = ContreeTestClient() + tc.fake.responses.extend(responses) + + CLIENT.set(tc) + FORMATTER.set(JSONFormatter()) + SESSION_STORE.set(store) + ctx = copy_context() + + with ( + patch( + "contree_cli.cli.run._stream_events_until_close", + side_effect=BrokenPipeError, + ), + patch("contree_cli.cli.run.sys.stdin", _tty_stdin()), + # `cmd_run` reopens stdout to /dev/null on BrokenPipeError; + # short-circuit that so pytest keeps its capture intact. + patch("contree_cli.cli.run.os.dup2"), + patch( + "contree_cli.cli.run.os.open", + return_value=os.open(os.devnull, os.O_RDONLY), + ), + patch("contree_cli.cli.run.os.close"), + pytest.raises(SystemExit) as exc_info, + ): + ctx.run(cmd_run, args) + assert exc_info.value.code == 141 + return tc + + def test_broken_pipe_cancels_operation(self, session_store): + tc = self._run_broken_pipe( + [_spawn_response(), _api_response({}, status=202)], + session_store, + ) + methods = [r.method for r in tc.fake.requests] + assert "DELETE" in methods + delete_req = next(r for r in tc.fake.requests if r.method == "DELETE") + assert "/v1/operations/op-1" in delete_req.path + + def test_broken_pipe_delete_failure_still_exits_141(self, session_store): + """Even if the DELETE fails, we still exit 141 rather than + re-raising BrokenPipeError.""" + self._run_broken_pipe( + [_spawn_response(), _api_response({"error": "not found"}, status=404)], + session_store, + ) + + # ── File upload ────────────────────────────────────────────────────────── @@ -552,13 +621,11 @@ def test_file_dedup_non_404_raises(self, contree_client, session_store, tmp_path ) args = _default_args(file=[mf]) - # Client retries 500s (RETRY_DELAYS has 7 entries -> 8 attempts) - responses = [ - _api_response({"error": "server error"}, status=500) for _ in range(8) - ] + # 403 is non-retryable, so the caller sees it on the first hit. + responses = [_api_response({"error": "forbidden"}, status=403)] with pytest.raises(ApiError) as exc_info: _run_cmd(contree_client, args, responses, store=session_store) - assert exc_info.value.status == 500 + assert exc_info.value.status == 403 def test_local_file_cache_skips_api_file_lookup( self, contree_client, session_store, tmp_path @@ -1968,3 +2035,433 @@ def test_use_without_command_still_sets_image( ] _run_cmd(contree_client, args, responses, store=session_store) assert session_store.session is not None + + +# --------------------------------------------------------------------------- +# SSE streaming (`_stream_events_until_close` and `_build_op_from_summary`) +# --------------------------------------------------------------------------- + + +class _SSEResponse: + """Minimal HTTPResponse-shaped object backed by an in-memory buffer. + + `_stream_events_until_close` only touches `.readline()` (via + `iter_sse_events`) and `.close()`; nothing else is needed. + """ + + def __init__(self, body: bytes | str) -> None: + payload = body.encode("utf-8") if isinstance(body, str) else body + self.buf = io.BytesIO(payload) + self.closed = False + + def readline(self, size: int = -1) -> bytes: + return self.buf.readline() + + def close(self) -> None: + self.closed = True + + +class _JSONResponse: + """Minimal HTTPResponse-shape for the streamer's between-attempt + ``GET /operations/{uuid}`` terminal check. Only ``.read()`` and + ``.close()`` are used.""" + + def __init__(self, payload: dict) -> None: + self.payload = json.dumps(payload).encode("utf-8") + self.closed = False + + def read(self, size: int = -1) -> bytes: + return self.payload + + def close(self) -> None: + self.closed = True + + +class _StubStreamClient: + """Stand-in for `ContreeClient` that queues responses for both the + SSE `follow=1` endpoint and the between-attempt terminal-status + GET. + + `responses` feeds `GET /events?follow=1` calls. `get_responses` + feeds `GET /operations/{uuid}` terminal checks; when it's empty + the stub falls back to a non-terminal op (`status=EXECUTING`) so + tests that don't care about the check don't have to prime it. + + A queued item may be an `_SSEResponse` / `_JSONResponse` (served + as-is), an `Exception` subclass or instance (raised), or `None` + (equivalent to an empty response). All calls are recorded so + tests can inspect them via `.calls`, `.sse_calls`, `.get_calls`. + """ + + NON_TERMINAL: ClassVar[dict[str, str]] = {"status": "EXECUTING"} + + def __init__( + self, + responses: list, + get_responses: list | None = None, + ) -> None: + self.responses = list(responses) + self.get_responses = list(get_responses or []) + self.calls: list[tuple[str, str, dict[str, str] | None]] = [] + + @property + def sse_calls(self) -> list[tuple[str, str, dict[str, str] | None]]: + return [c for c in self.calls if "/events" in c[1]] + + @property + def get_calls(self) -> list[tuple[str, str, dict[str, str] | None]]: + return [c for c in self.calls if "/events" not in c[1]] + + def request( + self, + method: str, + path: str, + *, + headers: dict[str, str] | None = None, + **_: object, + ): + self.calls.append((method, path, headers)) + is_sse = "/events" in path + if is_sse: + item = self.responses.pop(0) + elif self.get_responses: + item = self.get_responses.pop(0) + else: + return _JSONResponse(self.NON_TERMINAL) + if isinstance(item, type) and issubclass(item, BaseException): + raise item("stub") + if isinstance(item, BaseException): + raise item + if item is None: + return _SSEResponse(b"") if is_sse else _JSONResponse(self.NON_TERMINAL) + return item + + +def _frame(*, id: int, event: str, data: dict | str) -> str: + """Build one SSE frame. Mirrors the API contract by embedding `id` + inside the JSON `data:` payload as well as the SSE `id:` line + (the CLI reads the id from the JSON, not the SSE header).""" + if isinstance(data, dict) and "id" not in data: + data = {"id": id, **data} + data_str = data if isinstance(data, str) else json.dumps(data) + return f"id: {id}\nevent: {event}\ndata: {data_str}\n\n" + + +class TestStreamEventsUntilClose: + def test_url_uses_follow_1(self, session_store): + """Verifies the endpoint spelling matches the OpenAPI spec (`?follow=1`).""" + completion = _frame( + id=1, + event="completion", + data={"type": "completion", "data": {"status": "SUCCESS"}}, + ) + client = _StubStreamClient([_SSEResponse(completion)]) + _stream_events_until_close(client, "op-1", DefaultFormatter()) + assert client.calls[0][0] == "GET" + assert client.calls[0][1] == "/v1/operations/op-1/events?follow=1" + + def test_first_call_has_no_last_event_id_header(self, session_store): + client = _StubStreamClient( + [ + _SSEResponse( + _frame( + id=1, + event="completion", + data={"type": "completion", "data": {}}, + ) + ) + ] + ) + _stream_events_until_close(client, "op-x", DefaultFormatter()) + assert client.calls[0][2] is None + + def test_stdout_streamed_live_for_default_formatter(self, capsys): + chunk = _frame( + id=1, + event="stdout", + data={ + "type": "stdout", + "spid": 1, + "data": {"value": "hello", "encoding": "ascii"}, + }, + ) + completion = _frame( + id=2, + event="completion", + data={"type": "completion", "data": {"status": "SUCCESS"}}, + ) + client = _StubStreamClient([_SSEResponse(chunk + completion)]) + summary = _stream_events_until_close(client, "op-1", DefaultFormatter()) + out = capsys.readouterr() + assert out.out == "hello" + assert bytes(summary.stdout) == b"hello" + + def test_stderr_streamed_live_for_default_formatter(self, capsys): + chunk = _frame( + id=1, + event="stderr", + data={ + "type": "stderr", + "spid": 1, + "data": {"value": "oops\n", "encoding": "ascii"}, + }, + ) + completion = _frame( + id=2, + event="completion", + data={"type": "completion", "data": {"status": "FAILED"}}, + ) + client = _StubStreamClient([_SSEResponse(chunk + completion)]) + summary = _stream_events_until_close(client, "op-1", DefaultFormatter()) + out = capsys.readouterr() + assert out.err == "oops\n" + assert bytes(summary.stderr) == b"oops\n" + + def test_json_formatter_accumulates_without_printing(self, capsys): + chunk = _frame( + id=1, + event="stdout", + data={ + "type": "stdout", + "spid": 1, + "data": {"value": "hi", "encoding": "ascii"}, + }, + ) + completion = _frame( + id=2, + event="completion", + data={"type": "completion", "data": {"status": "SUCCESS"}}, + ) + client = _StubStreamClient([_SSEResponse(chunk + completion)]) + summary = _stream_events_until_close(client, "op-1", JSONFormatter()) + out = capsys.readouterr() + assert out.out == "" + assert bytes(summary.stdout) == b"hi" + + def test_base64_chunk_decoded(self, capsys): + payload = base64.b64encode(b"\x00\xff bin").decode("ascii") + chunk = _frame( + id=1, + event="stdout", + data={ + "type": "stdout", + "spid": 1, + "data": {"value": payload, "encoding": "base64"}, + }, + ) + completion = _frame( + id=2, + event="completion", + data={"type": "completion", "data": {"status": "SUCCESS"}}, + ) + client = _StubStreamClient([_SSEResponse(chunk + completion)]) + summary = _stream_events_until_close(client, "op-1", JSONFormatter()) + assert bytes(summary.stdout) == b"\x00\xff bin" + + def test_exit_event_for_spid_1_stored(self): + exit_frame = _frame( + id=1, + event="exit", + data={"type": "exit", "spid": 1, "data": {"code": 0, "timed_out": False}}, + ) + completion = _frame( + id=2, + event="completion", + data={"type": "completion", "data": {"status": "SUCCESS"}}, + ) + client = _StubStreamClient([_SSEResponse(exit_frame + completion)]) + summary = _stream_events_until_close(client, "op-1", DefaultFormatter()) + assert summary.exit_event is not None + assert summary.exit_event["data"]["code"] == 0 + + def test_exit_event_for_non_main_spid_ignored(self): + """Only spid=1 drives CLI exit code; child spid exits stay unrecorded.""" + exit_frame = _frame( + id=1, + event="exit", + data={"type": "exit", "spid": 2, "data": {"code": 3, "timed_out": False}}, + ) + completion = _frame( + id=2, + event="completion", + data={"type": "completion", "data": {"status": "SUCCESS"}}, + ) + client = _StubStreamClient([_SSEResponse(exit_frame + completion)]) + summary = _stream_events_until_close(client, "op-1", DefaultFormatter()) + assert summary.exit_event is None + + def test_completion_frame_breaks_loop(self): + completion = _frame( + id=1, + event="completion", + data={"type": "completion", "data": {"status": "SUCCESS"}}, + ) + client = _StubStreamClient([_SSEResponse(completion)]) + summary = _stream_events_until_close(client, "op-1", DefaultFormatter()) + assert summary.completion is not None + assert summary.completion["data"]["status"] == "SUCCESS" + + def test_stream_ends_without_completion_falls_back_to_terminal_get(self): + """SSE closes cleanly without a `completion` frame — the + between-attempt GET check detects the op is terminal and parks + it on `summary.fallback_op` so the caller doesn't need to GET + again.""" + op = {"uuid": "op-1", "status": "SUCCESS", "result": {"image": None}} + client = _StubStreamClient( + [_SSEResponse(b"")], get_responses=[_JSONResponse(op)] + ) + summary = _stream_events_until_close(client, "op-1", DefaultFormatter()) + assert summary.completion is None + assert summary.fallback_op == op + assert bytes(summary.stdout) == b"" + + def test_sse_connect_errors_then_terminal_via_get(self): + """SSE keeps failing to connect while the op runs; once the + between-attempt GET reports terminal status the streamer + stops retrying and returns without ever seeing a completion + frame.""" + op = {"uuid": "op-1", "status": "SUCCESS"} + client = _StubStreamClient( + [ApiError(500, "srv", "boom")] * 3, + get_responses=[ + _JSONResponse({"status": "EXECUTING"}), + _JSONResponse({"status": "EXECUTING"}), + _JSONResponse(op), + ], + ) + with patch("contree_cli.cli.run.time.sleep"): + summary = _stream_events_until_close(client, "op-1", DefaultFormatter()) + assert summary.completion is None + assert summary.fallback_op == op + assert len(client.sse_calls) == 3 + assert len(client.get_calls) == 3 + + def test_sse_error_frame_triggers_reconnect_with_last_event_id(self): + first = ( + _frame( + id=1, + event="stdout", + data={ + "type": "stdout", + "spid": 1, + "data": {"value": "a", "encoding": "ascii"}, + }, + ) + + "event: sse_error\ndata: boom\n\n" + ) + completion = _frame( + id=2, + event="completion", + data={"type": "completion", "data": {"status": "SUCCESS"}}, + ) + client = _StubStreamClient([_SSEResponse(first), _SSEResponse(completion)]) + with patch("contree_cli.cli.run.time.sleep"): + summary = _stream_events_until_close(client, "op-1", DefaultFormatter()) + assert summary.completion is not None + assert len(client.sse_calls) == 2 + assert client.sse_calls[1][2] == {"Last-Event-Id": "1"} + + def test_broken_pipe_from_stdout_propagates(self, monkeypatch): + """`BrokenPipeError` from local stdio write must propagate + unchanged — retrying can't fix a closed local pipe and it + would be misinterpreted as a remote network error otherwise.""" + chunk = _frame( + id=1, + event="stdout", + data={ + "type": "stdout", + "spid": 1, + "data": {"value": "hi", "encoding": "ascii"}, + }, + ) + client = _StubStreamClient([_SSEResponse(chunk)]) + + def raise_broken_pipe(*_args: object, **_kw: object) -> int: + raise BrokenPipeError + + monkeypatch.setattr("sys.stdout.buffer.write", raise_broken_pipe) + with pytest.raises(BrokenPipeError): + _stream_events_until_close(client, "op-1", DefaultFormatter()) + # Only the initial SSE attempt is made; no retry, no terminal-check + # GET (BrokenPipeError bypasses the retry path entirely). + assert len(client.sse_calls) == 1 + assert len(client.get_calls) == 0 + + +class TestBuildOpFromSummary: + def _completion(self, **overrides) -> dict: + base = { + "type": "completion", + "data": { + "status": "SUCCESS", + "error": None, + "duration_ms": 1500, + "image_size_bytes": 4096, + "result_image_uuid": IMG_NEW, + }, + } + base["data"].update(overrides) + return base + + def test_shape_carries_status_and_uuid(self): + summary = TerminalSummary(completion=self._completion()) + op = _build_op_from_summary("op-1", summary) + assert op["uuid"] == "op-1" + assert op["kind"] == "instance" + assert op["status"] == "SUCCESS" + + def test_duration_ms_converted_to_seconds(self): + summary = TerminalSummary(completion=self._completion(duration_ms=2500)) + op = _build_op_from_summary("op-1", summary) + assert op["duration"] == 2.5 + + def test_duration_missing_falls_back_to_zero(self): + summary = TerminalSummary( + completion={"type": "completion", "data": {"status": "SUCCESS"}} + ) + op = _build_op_from_summary("op-1", summary) + assert op["duration"] == 0.0 + + def test_result_image_uuid_propagates(self): + summary = TerminalSummary(completion=self._completion()) + op = _build_op_from_summary("op-1", summary) + assert op["result_image_uuid"] == IMG_NEW + assert op["result"] == {"image": IMG_NEW, "tag": None} + + def test_exit_event_drives_state(self): + summary = TerminalSummary( + completion=self._completion(), + exit_event={ + "type": "exit", + "spid": 1, + "data": {"code": 42, "timed_out": True}, + }, + ) + op = _build_op_from_summary("op-1", summary) + state = op["metadata"]["result"]["state"] + assert state == {"exit_code": 42, "timed_out": True} + + def test_state_is_none_without_exit_event(self): + summary = TerminalSummary(completion=self._completion()) + op = _build_op_from_summary("op-1", summary) + assert op["metadata"]["result"]["state"] is None + + def test_stdout_stderr_reassembled_from_bytearrays(self): + summary = TerminalSummary( + completion=self._completion(), + stdout=bytearray(b"hello\n"), + stderr=bytearray(b"warn\n"), + ) + op = _build_op_from_summary("op-1", summary) + result = op["metadata"]["result"] + assert result["stdout"]["value"] == "hello\n" + assert result["stderr"]["value"] == "warn\n" + assert result["stdout"]["truncated"] is False + + def test_error_field_passthrough(self): + summary = TerminalSummary( + completion=self._completion(status="FAILED", error="boom") + ) + op = _build_op_from_summary("op-1", summary) + assert op["status"] == "FAILED" + assert op["error"] == "boom" diff --git a/tests/test_url_fetch.py b/tests/test_url_fetch.py index d72d96b..d65d562 100644 --- a/tests/test_url_fetch.py +++ b/tests/test_url_fetch.py @@ -8,7 +8,6 @@ import contree_cli.docker.url_fetch as url_fetch from contree_cli.docker.url_fetch import ( - HashingReader, fetch_and_upload, is_url, url_basename, @@ -62,23 +61,6 @@ def test_validators_match_returns_false_when_no_validators(self): assert validators_match({}, {}) is False -class TestHashingReader: - def test_hashes_and_counts(self): - src = FakeStream(b"hello world") - r = HashingReader(src) - chunks = [r.read(5), r.read(5), r.read(5)] - assert b"".join(chunks) == b"hello world" - assert r.bytes_read == 11 - import hashlib - - assert r.hasher.hexdigest() == hashlib.sha256(b"hello world").hexdigest() - - def test_has_no_seek_attribute(self): - """ContreeClient.request relies on absent .seek to skip retry-rewind.""" - r = HashingReader(FakeStream(b"x")) - assert not hasattr(r, "seek") - - class TestFetchAndUpload: URL = "https://example.com/pkg.tgz"