contree-client 0.2.0__tar.gz → 0.2.1.dev0__tar.gz

{contree_client-0.2.0 → contree_client-0.2.1.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: contree-client
-Version: 0.2.0
+Version: 0.2.1.dev0
 Summary: Python client for the Contree API, generated from the OpenAPI specification
 Author: Dmitry Orlov
 Author-email: Dmitry Orlov <mosquito@nebius.com>

{contree_client-0.2.0 → contree_client-0.2.1.dev0}/contree_client/aiohttp.py

@@ -35,6 +35,7 @@ class ContreeAsyncClient(base.ContreeAsyncClient):
     retryable_errors = (
         aiohttp.ClientConnectionError,
         aiohttp.ServerConnectionError,
+        aiohttp.ClientPayloadError,
     )
     nonretryable_errors = (TimeoutError, asyncio.TimeoutError)
 
@@ -126,7 +127,12 @@ class ContreeAsyncClient(base.ContreeAsyncClient):
         ) as response:
             self.log.debug("%s %s -> %d (stream)", spec.method, url, response.status)
             if not 200 <= response.status < 300:
-                body = await response.read()
+                try:
+                    body = await response.read()
+                except self.retryable_errors:
+                    # The status is authoritative even when its optional
+                    # diagnostic body is interrupted.
+                    body = b""
                 # auto_decompress=False applies to the payload only:
                 # the error body must still be decoded, or the parsed
                 # server message is lost

{contree_client-0.2.0 → contree_client-0.2.1.dev0}/contree_client/base.py

@@ -1625,6 +1625,8 @@ class ContreeSyncClient(ContreeClientBase, ABC):
         max_count: int | None = None,
         max_total: int | None = None,
         case: Literal["sensitive", "insensitive", "smart"] | None = None,
+        before: int | None = None,
+        after: int | None = None,
     ) -> GrepResult:
         """Search file contents in image.
 
@@ -1637,20 +1639,39 @@ class ContreeSyncClient(ContreeClientBase, ABC):
         (containing a NUL byte) are skipped. Rust regex syntax, linear-time
         matching.
 
+        Results are a flat, file-ordered list of lines. ``before``/``after``
+        request extra
+        ``type: context`` lines around each ``type: match`` line, mirroring
+        ripgrep's
+        --json message type field.
+
+        ``path`` may be repeated to search multiple roots in one call, the same
+        way
+        ``pattern``/``glob`` already can; results are concatenated in the given
+        order
+        and share the ``max_total`` budget across all of them.
+
         Args:
             image_uuid: The UUID of the image to inspect Example:
                 ``12345678-9abc-baba-deda-0123456789ab``.
             pattern: Regex pattern to search for. May be repeated (max 16);
                 patterns are OR-ed. Max 1024 characters each. Example:
                 ``^PermitRootLogin``.
-            path: File or directory to search inside the image. Defaults to
-                the rootfs root. Example: ``/etc``.
+            path: File or directory to search inside the image. May be
+                repeated (max 16) to search multiple roots; results are
+                concatenated in the given order, sharing max_total across
+                all paths. Defaults to the rootfs root. Example: ``/etc``.
             glob: Glob to filter files (ripgrep -g semantics, "!" negates).
                 May be repeated. Example: ``*.conf``.
             max_count: Maximum matches per file (hard cap 10000).
             max_total: Maximum total matches across all files (default 1000,
                 hard cap 10000).
             case: Case sensitivity mode.
+            before: Lines of context to include before each match (ripgrep
+                -B semantics, hard cap 50). For symmetric context (grep -C
+                n), pass both before and after yourself.
+            after: Lines of context to include after each match (ripgrep -A
+                semantics, hard cap 50).
         """
         spec = operations.build_inspect_image_grep(
             image_uuid=image_uuid,
@@ -1660,6 +1681,8 @@ class ContreeSyncClient(ContreeClientBase, ABC):
             max_count=max_count,
             max_total=max_total,
             case=case,
+            before=before,
+            after=after,
         )
         self.log_request(spec)
         response = self.call(spec)
@@ -1803,9 +1826,26 @@ class ContreeAsyncClient(ContreeClientBase, ABC):
         ``OperationResponse``. Raises :class:`TimeoutError` when
         *timeout* seconds elapse first.
         """
-        async for event in self.follow_operation_events(operation_id, timeout=timeout):
-            self.log.debug("wait_operation: event %s %s", event.id, event.type)
-        return await self.get_operation_status(operation_id)
+
+        async def wait_and_fetch() -> OperationResponse:
+            async for event in self.follow_operation_events(operation_id):
+                self.log.debug("wait_operation: event %s %s", event.id, event.type)
+            return await self.get_operation_status(operation_id)
+
+        if timeout is None:
+            return await wait_and_fetch()
+        task = asyncio.create_task(wait_and_fetch())
+        try:
+            done, _pending = await asyncio.wait((task,), timeout=timeout)
+            if task in done:
+                return task.result()
+        finally:
+            if not task.done():
+                task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+        raise TimeoutError(
+            f"operation {operation_id} did not complete within {timeout}s"
+        )
 
     async def follow_operation_events(
         self,
@@ -1839,6 +1879,21 @@ class ContreeAsyncClient(ContreeClientBase, ABC):
                     f" within {timeout}s"
                 )
 
+        async def operation_terminal_before_deadline() -> bool:
+            if deadline is None:
+                return await self.operation_terminal(operation_id)
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                check_deadline()
+            try:
+                return await asyncio.wait_for(
+                    self.operation_terminal(operation_id),
+                    timeout=remaining,
+                )
+            except asyncio.TimeoutError:
+                check_deadline()
+                raise
+
         while True:
             check_deadline()
             events_before = last_id
@@ -1871,7 +1926,7 @@ class ContreeAsyncClient(ContreeClientBase, ABC):
                 retryable = exc.status in (410, 425) or 500 <= exc.status < 600
                 if not retryable:
                     raise
-                if await self.operation_terminal(operation_id):
+                if await operation_terminal_before_deadline():
                     return
                 self.log.warning("stream connect failed (%d): %s", exc.status, exc)
                 retry_after = exc.retry_after
@@ -1883,14 +1938,18 @@ class ContreeAsyncClient(ContreeClientBase, ABC):
                 continue
             except self.retryable_errors as exc:
                 if isinstance(exc, self.nonretryable_errors):
+                    check_deadline()
                     raise
                 self.log.warning("stream broken (last_id=%s): %s", last_id, exc)
             # the stream ended or broke without a completion frame:
             # the retry must not outlive the operation itself
-            if await self.operation_terminal(operation_id):
+            if await operation_terminal_before_deadline():
                 return
             if last_id == events_before:
-                await asyncio.sleep(TIGHT_LOOP_FLOOR)
+                delay = TIGHT_LOOP_FLOOR
+                if deadline is not None:
+                    delay = min(delay, max(0.0, deadline - time.monotonic()))
+                await asyncio.sleep(delay)
 
     async def resolve_image(self, ref: str) -> str:
         """Resolve an image reference to a UUID.
@@ -3097,6 +3156,8 @@ class ContreeAsyncClient(ContreeClientBase, ABC):
         max_count: int | None = None,
         max_total: int | None = None,
         case: Literal["sensitive", "insensitive", "smart"] | None = None,
+        before: int | None = None,
+        after: int | None = None,
     ) -> GrepResult:
         """Search file contents in image.
 
@@ -3109,20 +3170,39 @@ class ContreeAsyncClient(ContreeClientBase, ABC):
         (containing a NUL byte) are skipped. Rust regex syntax, linear-time
         matching.
 
+        Results are a flat, file-ordered list of lines. ``before``/``after``
+        request extra
+        ``type: context`` lines around each ``type: match`` line, mirroring
+        ripgrep's
+        --json message type field.
+
+        ``path`` may be repeated to search multiple roots in one call, the same
+        way
+        ``pattern``/``glob`` already can; results are concatenated in the given
+        order
+        and share the ``max_total`` budget across all of them.
+
         Args:
             image_uuid: The UUID of the image to inspect Example:
                 ``12345678-9abc-baba-deda-0123456789ab``.
             pattern: Regex pattern to search for. May be repeated (max 16);
                 patterns are OR-ed. Max 1024 characters each. Example:
                 ``^PermitRootLogin``.
-            path: File or directory to search inside the image. Defaults to
-                the rootfs root. Example: ``/etc``.
+            path: File or directory to search inside the image. May be
+                repeated (max 16) to search multiple roots; results are
+                concatenated in the given order, sharing max_total across
+                all paths. Defaults to the rootfs root. Example: ``/etc``.
             glob: Glob to filter files (ripgrep -g semantics, "!" negates).
                 May be repeated. Example: ``*.conf``.
             max_count: Maximum matches per file (hard cap 10000).
             max_total: Maximum total matches across all files (default 1000,
                 hard cap 10000).
             case: Case sensitivity mode.
+            before: Lines of context to include before each match (ripgrep
+                -B semantics, hard cap 50). For symmetric context (grep -C
+                n), pass both before and after yourself.
+            after: Lines of context to include after each match (ripgrep -A
+                semantics, hard cap 50).
         """
         spec = operations.build_inspect_image_grep(
             image_uuid=image_uuid,
@@ -3132,6 +3212,8 @@ class ContreeAsyncClient(ContreeClientBase, ABC):
             max_count=max_count,
             max_total=max_total,
             case=case,
+            before=before,
+            after=after,
         )
         self.log_request(spec)
         response = await self.call(spec)

{contree_client-0.2.0 → contree_client-0.2.1.dev0}/contree_client/models.py

@@ -285,7 +285,18 @@ class GrepMatch(ContreeModel):
             "example": 8192,
         }
     )
-    submatches: list[GrepSubmatch]
+    submatches: list[GrepSubmatch] = field(
+        metadata={"description": "Empty for context lines (they carry no submatches)"}
+    )
+    type: Literal["match", "context"] = field(
+        metadata={
+            "description": (
+                'Mirrors ripgrep\'s --json message type: "match" is an actual match, '
+                '"context" is a\nbefore/after context line requested via before/after.'
+            ),
+            "example": "match",
+        }
+    )
 
     @classmethod
     def parse_fields(cls, data: dict[str, Any]) -> dict[str, Any]:
@@ -296,6 +307,7 @@ class GrepMatch(ContreeModel):
             "line_text": data["line_text"],
             "line_bytes": data["line_bytes"],
             "submatches": [GrepSubmatch.from_dict(item) for item in data["submatches"]],
+            "type": data["type"],
         }
 
 
@@ -303,7 +315,11 @@ class GrepMatch(ContreeModel):
 class GrepResult(ContreeModel):
     path: str = field(
         metadata={
-            "description": "The path inside the image that was searched",
+            "description": (
+                "The first path inside the image that was searched. When `path` was "
+                "repeated,\nsee each row's own `path` for which root it actually came "
+                "from."
+            ),
             "example": "/etc",
         }
     )
@@ -313,7 +329,14 @@ class GrepResult(ContreeModel):
             "example": ["^PermitRootLogin"],
         }
     )
-    matches: list[GrepMatch]
+    matches: list[GrepMatch] = field(
+        metadata={
+            "description": (
+                "File-ordered lines; each is a match or (when before/after are set) "
+                "surrounding context"
+            )
+        }
+    )
     truncated: bool = field(
         metadata={
             "description": (

{contree_client-0.2.0 → contree_client-0.2.1.dev0}/contree_client/operations.py

@@ -598,6 +598,8 @@ def build_inspect_image_grep(
     max_count: int | None = None,
     max_total: int | None = None,
     case: Literal["sensitive", "insensitive", "smart"] | None = None,
+    before: int | None = None,
+    after: int | None = None,
 ) -> RequestSpec:
     """Build the request for `GET /inspect/{image_uuid}/grep`."""
     query: dict[str, str] = {}
@@ -612,6 +614,10 @@ def build_inspect_image_grep(
         query["max_total"] = str(max_total)
     if case is not None:
         query["case"] = case
+    if before is not None:
+        query["before"] = str(before)
+    if after is not None:
+        query["after"] = str(after)
     request_path = f"/inspect/{quote_path(image_uuid)}/grep"
     return RequestSpec(method="GET", path=request_path, idempotent=True, query=query)
 

{contree_client-0.2.0 → contree_client-0.2.1.dev0}/contree_client/spec_info.py

@@ -1571,6 +1571,14 @@ Generated by codegen from the OpenAPI spec - do not edit.
             Symlinks are never followed during traversal, hidden files are searched,
             `.gitignore`-style files inside the image are ignored, and binary files
             (containing a NUL byte) are skipped. Rust regex syntax, linear-time matching.
+
+            Results are a flat, file-ordered list of lines. `before`/`after` request extra
+            `type: context` lines around each `type: match` line, mirroring ripgrep's
+            --json message type field.
+
+            `path` may be repeated to search multiple roots in one call, the same way
+            `pattern`/`glob` already can; results are concatenated in the given order
+            and share the `max_total` budget across all of them.
           operationId: inspectImageGrep
           tags:
             - inspect
@@ -1593,7 +1601,10 @@ Generated by codegen from the OpenAPI spec - do not edit.
               in: query
               required: false
               allowReserved: true
-              description: File or directory to search inside the image. Defaults to the rootfs root.
+              description: |
+                File or directory to search inside the image. May be repeated (max 16) to search
+                multiple roots; results are concatenated in the given order, sharing max_total
+                across all paths. Defaults to the rootfs root.
               schema:
                 type: string
                 default: "/"
@@ -1631,6 +1642,26 @@ Generated by codegen from the OpenAPI spec - do not edit.
                 type: string
                 enum: [sensitive, insensitive, smart]
                 default: sensitive
+            - name: before
+              in: query
+              required: false
+              description: |
+                Lines of context to include before each match (ripgrep -B semantics, hard cap 50).
+                For symmetric context (grep -C n), pass both before and after yourself.
+              schema:
+                type: integer
+                minimum: 0
+                maximum: 50
+                default: 0
+            - name: after
+              in: query
+              required: false
+              description: Lines of context to include after each match (ripgrep -A semantics, hard cap 50).
+              schema:
+                type: integer
+                minimum: 0
+                maximum: 50
+                default: 0
           responses:
             '200':
               description: OK
@@ -1928,6 +1959,7 @@ Generated by codegen from the OpenAPI spec - do not edit.
             - line_text
             - line_bytes
             - submatches
+            - type
           properties:
             path:
               type: string
@@ -1951,8 +1983,16 @@ Generated by codegen from the OpenAPI spec - do not edit.
               example: 8192
             submatches:
               type: array
+              description: Empty for context lines (they carry no submatches)
               items:
                 $ref: '#/components/schemas/GrepSubmatch'
+            type:
+              type: string
+              enum: [match, context]
+              description: |
+                Mirrors ripgrep's --json message type: "match" is an actual match, "context" is a
+                before/after context line requested via before/after.
+              example: match
 
         GrepResult:
           type: object
@@ -1964,7 +2004,9 @@ Generated by codegen from the OpenAPI spec - do not edit.
           properties:
             path:
               type: string
-              description: The path inside the image that was searched
+              description: |
+                The first path inside the image that was searched. When `path` was repeated,
+                see each row's own `path` for which root it actually came from.
               example: "/etc"
             patterns:
               type: array
@@ -1974,6 +2016,7 @@ Generated by codegen from the OpenAPI spec - do not edit.
               example: ["^PermitRootLogin"]
             matches:
               type: array
+              description: File-ordered lines; each is a match or (when before/after are set) surrounding context
               items:
                 $ref: '#/components/schemas/GrepMatch'
             truncated:
@@ -3403,4 +3446,4 @@ DEFAULT_BASE_URL = "https://api.tokenfactory.nebius.com/sandboxes"
 
 # sha256 of the exact OpenAPI document this package was built from -
 # the build input provenance
-SPEC_SHA256 = "bf17dbf27d28bebe3e23b295835bbdb2927ab9de5e9ae581af675e90986acc56"
+SPEC_SHA256 = "2bfd239341e6f61ccba1aafb39cef7b476d56f424156860a204f089d7d79b7bf"

contree_client-0.2.1.dev0/pyproject.toml

@@ -0,0 +1,78 @@
+[build-system]
+requires = ["uv_build>=0.11.28,<0.12"]
+build-backend = "uv_build"
+
+[project]
+name = "contree-client"
+version = "0.2.1.dev0"
+description = "Python client for the Contree API, generated from the OpenAPI specification"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "Apache-2.0"
+license-files = ["LICENSE"]
+dependencies = []
+
+[[project.authors]]
+name = "Dmitry Orlov"
+email = "mosquito@nebius.com"
+
+[project.urls]
+Homepage = "https://contree.dev/"
+Documentation = "https://docs.tokenfactory.nebius.com/sandboxes/overview"
+
+[project.optional-dependencies]
+urllib3 = ["urllib3>=2.0"]
+requests = ["requests>=2.28"]
+aiohttp = ["aiohttp>=3.9"]
+httpx = ["httpx>=0.27"]
+dev = [
+    "contree-client[urllib3,requests,aiohttp,httpx]",
+    "pytest>=8",
+    "pytest-cov>=7.1.0",
+    "pytest-timeout>=2.3",
+    "markdown-pytest[async]>=0.3",
+    "types-requests",
+]
+
+[tool.uv.build-backend]
+module-root = ""
+namespace = true
+source-include = ["tests/**"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+markers = ["integration: live tests against the real Contree API; auto-skip without an active contree-cli profile"]
+
+[tool.ty.environment]
+python-version = "3.10"
+
+[tool.ty.src]
+respect-ignore-files = false
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.ruff.lint]
+select = [
+    "A",
+    "B",
+    "C4",
+    "E",
+    "F",
+    "FURB",
+    "I",
+    "ISC",
+    "PERF",
+    "PIE",
+    "PLC",
+    "PLE",
+    "RET",
+    "RSE",
+    "RUF",
+    "SIM",
+    "T20",
+    "UP",
+    "W",
+]
+ignore = ["ISC001"]

contree_client-0.2.0/pyproject.toml → contree_client-0.2.1.dev0/pyproject.toml.orig

@@ -6,7 +6,7 @@ build-backend = "uv_build"
 
 [project]
 name = "contree-client"
-version = "0.2.0"
+version = "0.2.1.dev0"
 description = "Python client for the Contree API, generated from the OpenAPI specification"
 readme = "README.md"
 requires-python = ">=3.10"

{contree_client-0.2.0 → contree_client-0.2.1.dev0}/tests/stub_server.py

@@ -32,11 +32,17 @@ RETRY_AFTER_OPERATION_UUID = "00000000-0000-0000-0000-000000000425"
 PENDING_OPERATION_UUID = "00000000-0000-0000-0000-0000000e4ec1"
 EXECUTING_OPERATION_UUID = "00000000-0000-0000-0000-0000000e4ec2"
 RECONNECT_OPERATION_UUID = "00000000-0000-0000-0000-00000000ec0e"
+PAYLOAD_INTERRUPTION_OPERATION_UUID = "00000000-0000-0000-0000-00000000badd"
+PAYLOAD_TIMEOUT_OPERATION_UUID = "00000000-0000-0000-0000-00000000dead"
+PAYLOAD_STALLED_STATUS_OPERATION_UUID = "00000000-0000-0000-0000-0000000057a1"
+PAYLOAD_FORBIDDEN_OPERATION_UUID = "00000000-0000-0000-0000-000000000403"
+PAYLOAD_FINAL_STATUS_OPERATION_UUID = "00000000-0000-0000-0000-00000000f1a1"
 CLOSING_OPERATION_UUID = "00000000-0000-0000-0000-0000000c105e"
 DROPPING_OPERATION_UUID = "00000000-0000-0000-0000-00000000d40b"
 KEEPALIVE_OPERATION_UUID = "00000000-0000-0000-0000-0000000cee9a"
 RESET_OPERATION_UUID = "00000000-0000-0000-0000-000000000e5e"
 SSE_HANG_SECONDS = 10.0
+STATUS_HANG_SECONDS = 2.0
 FILE_UUID = "a9165a5d-5c86-4bd8-8ee4-ae46c19cf45d"
 KNOWN_SHA256 = "a" * 64
 CREATED_AT = "2024-01-01T12:00:00+00:00"
@@ -285,6 +291,9 @@ class Reply:
     hang: float = 0.0
     # simulate a peer dying mid-transfer: the gzip trailer never comes
     truncate_gzip: bool = False
+    # simulate invalid HTTP chunk framing: close without the terminal
+    # zero-length chunk after delivering all declared chunks
+    truncate_chunked: bool = False
     # serve a normal keepalive response, then silently close the TCP
     # connection - the client only finds out when it tries to reuse it
     drop_after: bool = False
@@ -400,6 +409,106 @@ def route(request: Captured, attempts: collections.Counter[str]) -> Reply:
             stream_chunks=chunks,
         )
 
+    if (
+        path == f"/v1/operations/{PAYLOAD_INTERRUPTION_OPERATION_UUID}"
+        and method == "GET"
+    ):
+        status = "EXECUTING" if attempt == 1 else "SUCCESS"
+        return json_reply(
+            200,
+            {
+                **OPERATION_RESPONSE,
+                "uuid": PAYLOAD_INTERRUPTION_OPERATION_UUID,
+                "status": status,
+            },
+        )
+
+    if path == f"/v1/operations/{PAYLOAD_INTERRUPTION_OPERATION_UUID}/events":
+        if attempt == 1:
+            return Reply(
+                status=200,
+                content_type="text/event-stream",
+                stream_chunks=[sse_frame(EVENT_INIT), sse_frame(EVENT_SPAWN)],
+                truncate_chunked=True,
+            )
+        return Reply(
+            status=200,
+            content_type="text/event-stream",
+            stream_chunks=RESUMED_SSE_FRAMES,
+        )
+
+    if path == f"/v1/operations/{PAYLOAD_TIMEOUT_OPERATION_UUID}" and method == "GET":
+        return json_reply(
+            200,
+            {
+                **OPERATION_RESPONSE,
+                "uuid": PAYLOAD_TIMEOUT_OPERATION_UUID,
+                "status": "EXECUTING",
+            },
+        )
+
+    if path == f"/v1/operations/{PAYLOAD_TIMEOUT_OPERATION_UUID}/events":
+        chunks = [sse_frame(EVENT_INIT)] if attempt == 1 else [b": keepalive\n\n"]
+        return Reply(
+            status=200,
+            content_type="text/event-stream",
+            stream_chunks=chunks,
+            truncate_chunked=True,
+        )
+
+    if (
+        path == f"/v1/operations/{PAYLOAD_STALLED_STATUS_OPERATION_UUID}"
+        and method == "GET"
+    ):
+        time.sleep(STATUS_HANG_SECONDS)
+        return json_reply(
+            200,
+            {
+                **OPERATION_RESPONSE,
+                "uuid": PAYLOAD_STALLED_STATUS_OPERATION_UUID,
+                "status": "EXECUTING",
+            },
+        )
+
+    if path == f"/v1/operations/{PAYLOAD_STALLED_STATUS_OPERATION_UUID}/events":
+        return Reply(
+            status=200,
+            content_type="text/event-stream",
+            stream_chunks=[b": keepalive\n\n"],
+            truncate_chunked=True,
+        )
+
+    if path == f"/v1/operations/{PAYLOAD_FORBIDDEN_OPERATION_UUID}/events":
+        return Reply(
+            status=403,
+            content_type="application/json",
+            stream_chunks=[b'{"error":"forbidden","status":403}'],
+            truncate_chunked=True,
+        )
+
+    if (
+        path == f"/v1/operations/{PAYLOAD_FINAL_STATUS_OPERATION_UUID}"
+        and method == "GET"
+    ):
+        if attempt > 1:
+            time.sleep(STATUS_HANG_SECONDS)
+        return json_reply(
+            200,
+            {
+                **OPERATION_RESPONSE,
+                "uuid": PAYLOAD_FINAL_STATUS_OPERATION_UUID,
+                "status": "SUCCESS",
+            },
+        )
+
+    if path == f"/v1/operations/{PAYLOAD_FINAL_STATUS_OPERATION_UUID}/events":
+        return Reply(
+            status=200,
+            content_type="text/event-stream",
+            stream_chunks=[b": keepalive\n\n"],
+            truncate_chunked=True,
+        )
+
     if path == "/v1/images" and method == "GET":
         if query.get("limit") == ["0"]:
             return json_reply(
@@ -677,7 +786,8 @@ class Handler(BaseHTTPRequestHandler):
         self.send_header("Content-Length", str(len(payload)))
         self.end_headers()
         if method != "HEAD" and payload:
-            self.wfile.write(payload)
+            with contextlib.suppress(BrokenPipeError, ConnectionResetError, OSError):
+                self.wfile.write(payload)
         if reply.drop_after:
             self.close_connection = True
 
@@ -708,9 +818,10 @@ class Handler(BaseHTTPRequestHandler):
                 time.sleep(0.01)
             if reply.hang:
                 time.sleep(reply.hang)
-            if not reply.truncate_gzip:
-                self.write_chunk(compressor.flush(zlib.Z_FINISH))
-            self.write_chunk(b"")
+            if not reply.truncate_chunked:
+                if not reply.truncate_gzip:
+                    self.write_chunk(compressor.flush(zlib.Z_FINISH))
+                self.write_chunk(b"")
         except (BrokenPipeError, ConnectionResetError, OSError):
             pass
         finally:

contree_client-0.2.1.dev0/tests/test_aiohttp_payload_interruption.py

@@ -0,0 +1,234 @@
+"""aiohttp recovery from incomplete chunked operation-event responses."""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+import time
+from types import ModuleType
+from typing import Any
+
+import aiohttp
+import pytest
+
+from tests.stub_server import (
+    PAYLOAD_FINAL_STATUS_OPERATION_UUID,
+    PAYLOAD_FORBIDDEN_OPERATION_UUID,
+    PAYLOAD_INTERRUPTION_OPERATION_UUID,
+    PAYLOAD_STALLED_STATUS_OPERATION_UUID,
+    PAYLOAD_TIMEOUT_OPERATION_UUID,
+    StubServer,
+)
+
+
+def test_client_payload_error_is_retryable(
+    generated_package: ModuleType,
+) -> None:
+    module = importlib.import_module("contree_client.aiohttp")
+    error = aiohttp.ClientPayloadError("incomplete chunked response")
+
+    assert isinstance(error, module.ContreeAsyncClient.retryable_errors)
+    assert not isinstance(error, module.ContreeAsyncClient.nonretryable_errors)
+
+
+async def collect_events(
+    base_url: str,
+    operation_id: str = PAYLOAD_INTERRUPTION_OPERATION_UUID,
+    timeout: float = 1.0,
+) -> list[Any]:
+    module = importlib.import_module("contree_client.aiohttp")
+    async with module.ContreeAsyncClient("test-token", base_url=base_url) as client:
+        return [
+            event
+            async for event in client.follow_operation_events(
+                operation_id,
+                timeout=timeout,
+            )
+        ]
+
+
+def test_follow_operation_events_resumes_after_incomplete_chunked_body(
+    generated_package: ModuleType,
+    stub_server: StubServer,
+) -> None:
+    events = asyncio.run(collect_events(stub_server.base_url))
+
+    assert [event.id for event in events] == [0, 1, 2, 3]
+    event_requests = [
+        request
+        for request in stub_server.captured
+        if request.path
+        == f"/v1/operations/{PAYLOAD_INTERRUPTION_OPERATION_UUID}/events"
+    ]
+    assert len(event_requests) == 2
+    assert event_requests[1].headers["last-event-id"] == "1"
+
+
+def test_follow_operation_events_preserves_status_from_truncated_error_body(
+    generated_package: ModuleType,
+    stub_server: StubServer,
+    exceptions: ModuleType,
+) -> None:
+    with pytest.raises(exceptions.ForbiddenError) as caught:
+        asyncio.run(
+            collect_events(
+                stub_server.base_url,
+                PAYLOAD_FORBIDDEN_OPERATION_UUID,
+                timeout=0.2,
+            )
+        )
+
+    assert caught.value.status == 403
+    event_requests = [
+        request
+        for request in stub_server.captured
+        if request.path == f"/v1/operations/{PAYLOAD_FORBIDDEN_OPERATION_UUID}/events"
+    ]
+    assert len(event_requests) == 1
+
+
+async def wait_operation(base_url: str, operation_id: str, timeout: float) -> Any:
+    module = importlib.import_module("contree_client.aiohttp")
+    async with module.ContreeAsyncClient("test-token", base_url=base_url) as client:
+        return await client.wait_operation(operation_id, timeout=timeout)
+
+
+def test_wait_operation_resumes_after_incomplete_chunked_body(
+    generated_package: ModuleType,
+    stub_server: StubServer,
+) -> None:
+    operation = asyncio.run(
+        wait_operation(
+            stub_server.base_url,
+            PAYLOAD_INTERRUPTION_OPERATION_UUID,
+            1.0,
+        )
+    )
+
+    assert str(operation.status) == "SUCCESS"
+    event_requests = [
+        request
+        for request in stub_server.captured
+        if request.path
+        == f"/v1/operations/{PAYLOAD_INTERRUPTION_OPERATION_UUID}/events"
+    ]
+    assert len(event_requests) == 2
+    assert event_requests[1].headers["last-event-id"] == "1"
+
+
+def test_wait_operation_keeps_one_deadline_across_payload_interruptions(
+    generated_package: ModuleType,
+    stub_server: StubServer,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module = importlib.import_module("contree_client.base")
+    monkeypatch.setattr(module, "TIGHT_LOOP_FLOOR", 5.0)
+    started = time.monotonic()
+    with pytest.raises(TimeoutError, match=PAYLOAD_TIMEOUT_OPERATION_UUID):
+        asyncio.run(
+            wait_operation(
+                stub_server.base_url,
+                PAYLOAD_TIMEOUT_OPERATION_UUID,
+                0.2,
+            )
+        )
+
+    assert time.monotonic() - started < 1.0
+
+
+def test_wait_operation_deadline_bounds_terminal_status_probe(
+    generated_package: ModuleType,
+    stub_server: StubServer,
+) -> None:
+    started = time.monotonic()
+    with pytest.raises(TimeoutError, match=PAYLOAD_STALLED_STATUS_OPERATION_UUID):
+        asyncio.run(
+            wait_operation(
+                stub_server.base_url,
+                PAYLOAD_STALLED_STATUS_OPERATION_UUID,
+                0.2,
+            )
+        )
+
+    assert time.monotonic() - started < 1.0
+
+
+def test_wait_operation_deadline_bounds_final_status_fetch(
+    generated_package: ModuleType,
+    stub_server: StubServer,
+) -> None:
+    started = time.monotonic()
+    with pytest.raises(TimeoutError, match=PAYLOAD_FINAL_STATUS_OPERATION_UUID):
+        asyncio.run(
+            wait_operation(
+                stub_server.base_url,
+                PAYLOAD_FINAL_STATUS_OPERATION_UUID,
+                0.2,
+            )
+        )
+
+    assert time.monotonic() - started < 1.0
+
+
+async def follow_after_socket_timeout(delay: float, timeout: float) -> None:
+    module = importlib.import_module("contree_client.aiohttp")
+    client = module.ContreeAsyncClient("test-token", base_url="http://127.0.0.1")
+
+    async def interrupted_events(*args: Any, **kwargs: Any) -> Any:
+        await asyncio.sleep(delay)
+        raise aiohttp.SocketTimeoutError("read stalled")
+        yield  # pragma: no cover - makes this an async generator
+
+    client.iter_operation_events = interrupted_events
+    try:
+        async for _event in client.follow_operation_events(
+            PAYLOAD_TIMEOUT_OPERATION_UUID,
+            timeout=timeout,
+        ):
+            pass
+    finally:
+        await client.close()
+
+
+def test_follow_operation_events_normalizes_socket_timeout_at_deadline(
+    generated_package: ModuleType,
+) -> None:
+    with pytest.raises(TimeoutError, match=PAYLOAD_TIMEOUT_OPERATION_UUID):
+        asyncio.run(follow_after_socket_timeout(delay=0.05, timeout=0.01))
+
+
+def test_follow_operation_events_preserves_socket_timeout_before_deadline(
+    generated_package: ModuleType,
+) -> None:
+    with pytest.raises(aiohttp.SocketTimeoutError, match="read stalled"):
+        asyncio.run(follow_after_socket_timeout(delay=0.0, timeout=1.0))
+
+
+async def wait_after_socket_timeout(delay: float, timeout: float) -> None:
+    module = importlib.import_module("contree_client.aiohttp")
+    client = module.ContreeAsyncClient("test-token", base_url="http://127.0.0.1")
+
+    async def interrupted_events(*args: Any, **kwargs: Any) -> Any:
+        await asyncio.sleep(delay)
+        raise aiohttp.SocketTimeoutError("read stalled")
+        yield  # pragma: no cover - makes this an async generator
+
+    client.iter_operation_events = interrupted_events
+    try:
+        await client.wait_operation(PAYLOAD_TIMEOUT_OPERATION_UUID, timeout=timeout)
+    finally:
+        await client.close()
+
+
+def test_wait_operation_deadline_preempts_later_socket_timeout(
+    generated_package: ModuleType,
+) -> None:
+    with pytest.raises(TimeoutError, match=PAYLOAD_TIMEOUT_OPERATION_UUID):
+        asyncio.run(wait_after_socket_timeout(delay=0.05, timeout=0.01))
+
+
+def test_wait_operation_preserves_socket_timeout_before_its_deadline(
+    generated_package: ModuleType,
+) -> None:
+    with pytest.raises(aiohttp.SocketTimeoutError, match="read stalled"):
+        asyncio.run(wait_after_socket_timeout(delay=0.0, timeout=1.0))