pi-durable-python 1.3.0__py3-none-any.whl

pi_durable/tools/file_mutation_queue.py

@@ -0,0 +1,68 @@
+"""Serialization of ``write`` and ``edit`` mutations of one file within this process.
+
+Port of the v1.3.0 ``src/tools/file-mutation-queue.ts``. Files are keyed by the environment's
+:attr:`~pi_durable.env.FileSystem.id` and the canonical path.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from typing import TypeVar
+
+from ..context import Context
+from ..env import ExecutionEnv, get_or_throw
+from .path_utils import canonical_path
+
+T = TypeVar("T")
+
+#: Tail of the mutation chain of each file, keyed by file system id and canonical path.
+_queues: dict[str, asyncio.Future[None]] = {}
+
+
+async def _mutation_key(env: ExecutionEnv, path: str, context: Context) -> str:
+    absolute_path = get_or_throw(await env.absolute_path(path, context))
+    return f"{getattr(env, 'id', '')}\0{await canonical_path(env, absolute_path, context)}"
+
+
+async def with_file_mutation_queue(
+    env: ExecutionEnv,
+    path: str,
+    fn: Callable[[], Awaitable[T]],
+    context: Context,
+) -> T:
+    """Run ``fn`` after earlier mutations of the same file system and canonical path settle.
+
+    Other files, and other file systems, never wait. This is not a lock against ``bash`` or other
+    processes.
+    """
+
+    key = await _mutation_key(env, path, context)
+    loop = asyncio.get_running_loop()
+    previous = _queues.get(key)
+    if previous is None:
+        ready: asyncio.Future[None] = loop.create_future()
+        ready.set_result(None)
+        previous = ready
+    done: asyncio.Future[None] = loop.create_future()
+    # Take the slot without awaiting, so no other call can take it in between.
+    _queues[key] = done
+    entered = False
+    try:
+        # Shield the shared tail: cancelling this waiter must not cancel the mutation it waits on.
+        await asyncio.shield(previous)
+        entered = True
+        return await fn()
+    finally:
+        if not done.done():
+            done.set_result(None)
+        if _queues.get(key) is done:
+            # A cancelled waiter never mutates, so the next call must still wait for the mutation
+            # it was queued behind rather than proceed beside it.
+            if not entered and not previous.done():
+                _queues[key] = previous
+            else:
+                del _queues[key]
+
+
+__all__ = ["with_file_mutation_queue"]

pi_durable/tools/image.py

@@ -0,0 +1,248 @@
+"""Image type and header detection for the ``read`` tool.
+
+Port of the v1.3.0 ``src/tools/image.ts``. Only headers and, for PNG, the chunk headers up to the
+first ``acTL`` or ``IDAT`` are read; image pixels are never decoded here.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+from ..context import Context
+from ..env import BinaryReader, get_or_throw
+
+PNG_SIGNATURE = (0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A)
+#: Bytes every check except the APNG chunk walk needs: BMP reads up to offset 29.
+HEADER_BYTES = 32
+BLOCK_BYTES = 64 * 1024
+
+
+class ByteSource(Protocol):
+    """Positional reads of a file of ``size`` bytes."""
+
+    size: int
+
+    async def read(self, offset: int, length: int) -> bytes: ...
+
+
+async def detect_supported_image_mime_type_of(source: ByteSource) -> str | None:
+    """``detect_supported_image_mime_type`` of a whole file, reading only its header."""
+
+    header = await source.read(0, HEADER_BYTES)
+    if not _starts_with(header, PNG_SIGNATURE):
+        return detect_supported_image_mime_type(header)
+    if _is_png(header) and not await _is_animated_png_of(source):
+        return "image/png"
+    return None
+
+
+async def _is_animated_png_of(source: ByteSource) -> bool:
+    block = b""
+    block_start = 0
+
+    async def bytes_at(offset: int, length: int) -> bytes:
+        nonlocal block, block_start
+        if offset < block_start or offset + length > block_start + len(block):
+            block_start = offset
+            block = await source.read(offset, BLOCK_BYTES)
+        return block[offset - block_start : offset - block_start + length]
+
+    offset = len(PNG_SIGNATURE)
+    while offset + 8 <= source.size:
+        chunk_header = await bytes_at(offset, 8)
+        chunk_length = _read_uint32_be(chunk_header, 0)
+        if _starts_with_ascii(chunk_header, 4, "acTL"):
+            return True
+        if _starts_with_ascii(chunk_header, 4, "IDAT"):
+            return False
+        next_offset = offset + 8 + chunk_length + 4
+        if next_offset <= offset or next_offset > source.size:
+            return False
+        offset = next_offset
+    return False
+
+
+def detect_supported_image_mime_type(buffer: bytes) -> str | None:
+    """The MIME type of a supported image, or ``None``."""
+
+    if _starts_with(buffer, (0xFF, 0xD8, 0xFF)):
+        return None if _byte(buffer, 3) == 0xF7 else "image/jpeg"
+    if _starts_with(buffer, PNG_SIGNATURE):
+        return "image/png" if _is_png(buffer) and not _is_animated_png(buffer) else None
+    if _starts_with_ascii(buffer, 0, "GIF87a") or _starts_with_ascii(buffer, 0, "GIF89a"):
+        return "image/gif"
+    if _starts_with_ascii(buffer, 0, "RIFF") and _starts_with_ascii(buffer, 8, "WEBP"):
+        return "image/webp"
+    if _starts_with_ascii(buffer, 0, "BM") and _is_bmp(buffer):
+        return "image/bmp"
+    return None
+
+
+def _is_png(buffer: bytes) -> bool:
+    return (
+        len(buffer) >= 16
+        and _read_uint32_be(buffer, len(PNG_SIGNATURE)) == 13
+        and _starts_with_ascii(buffer, 12, "IHDR")
+    )
+
+
+def _is_animated_png(buffer: bytes) -> bool:
+    offset = len(PNG_SIGNATURE)
+    while offset + 8 <= len(buffer):
+        chunk_length = _read_uint32_be(buffer, offset)
+        chunk_type_offset = offset + 4
+        if _starts_with_ascii(buffer, chunk_type_offset, "acTL"):
+            return True
+        if _starts_with_ascii(buffer, chunk_type_offset, "IDAT"):
+            return False
+        next_offset = offset + 8 + chunk_length + 4
+        if next_offset <= offset or next_offset > len(buffer):
+            return False
+        offset = next_offset
+    return False
+
+
+def image_dimensions(bytes_: bytes, mime_type: str) -> dict[str, int] | None:
+    """The width and height an image declares in its header, without decoding it."""
+
+    length = len(bytes_)
+    if mime_type == "image/png":
+        if length < 24:
+            return None
+        return {"width": _read_uint32_be(bytes_, 16), "height": _read_uint32_be(bytes_, 20)}
+    if mime_type == "image/gif":
+        if length < 10:
+            return None
+        return {"width": _read_uint16_le(bytes_, 6), "height": _read_uint16_le(bytes_, 8)}
+    if mime_type == "image/webp":
+        if length < 30:
+            return None
+        if _starts_with_ascii(bytes_, 12, "VP8 "):
+            return {
+                "width": _read_uint16_le(bytes_, 26) & 0x3FFF,
+                "height": _read_uint16_le(bytes_, 28) & 0x3FFF,
+            }
+        if _starts_with_ascii(bytes_, 12, "VP8L"):
+            bits = _read_uint32_le(bytes_, 21)
+            return {
+                "width": (bits & 0x3FFF) + 1,
+                "height": ((bits >> 14) & 0x3FFF) + 1,
+            }
+        if _starts_with_ascii(bytes_, 12, "VP8X"):
+
+            def uint24(offset: int) -> int:
+                return _read_uint16_le(bytes_, offset) + (_byte(bytes_, offset + 2) << 16)
+
+            return {"width": uint24(24) + 1, "height": uint24(27) + 1}
+        return None
+    if mime_type == "image/jpeg":
+        offset = 2
+        while offset + 9 <= length:
+            if _byte(bytes_, offset) != 0xFF:
+                return None
+            marker = _byte(bytes_, offset + 1)
+            if marker == 0xFF:
+                offset += 1
+                continue
+            # SOF0 to SOF15, except DHT (C4), JPG (C8), and DAC (CC), which share the range.
+            if 0xC0 <= marker <= 0xCF and marker not in (0xC4, 0xC8, 0xCC):
+                return {
+                    "width": _read_uint16_be(bytes_, offset + 7),
+                    "height": _read_uint16_be(bytes_, offset + 5),
+                }
+            offset += 2 + _read_uint16_be(bytes_, offset + 2)
+        return None
+    return None
+
+
+def _is_bmp(buffer: bytes) -> bool:
+    if len(buffer) < 26:
+        return False
+    declared_file_size = _read_uint32_le(buffer, 2)
+    pixel_data_offset = _read_uint32_le(buffer, 10)
+    dib_header_size = _read_uint32_le(buffer, 14)
+    if declared_file_size != 0 and declared_file_size < 26:
+        return False
+    if pixel_data_offset < 14 + dib_header_size:
+        return False
+    if declared_file_size != 0 and pixel_data_offset >= declared_file_size:
+        return False
+
+    if dib_header_size == 12:
+        color_planes = _read_uint16_le(buffer, 22)
+        bits_per_pixel = _read_uint16_le(buffer, 24)
+    elif 40 <= dib_header_size <= 124:
+        if len(buffer) < 30:
+            return False
+        color_planes = _read_uint16_le(buffer, 26)
+        bits_per_pixel = _read_uint16_le(buffer, 28)
+    else:
+        return False
+    return color_planes == 1 and bits_per_pixel in (1, 4, 8, 16, 24, 32)
+
+
+def _byte(buffer: bytes, offset: int) -> int:
+    return buffer[offset] if offset < len(buffer) else 0
+
+
+def _read_uint16_le(buffer: bytes, offset: int) -> int:
+    return _byte(buffer, offset) + (_byte(buffer, offset + 1) << 8)
+
+
+def _read_uint16_be(buffer: bytes, offset: int) -> int:
+    return (_byte(buffer, offset) << 8) + _byte(buffer, offset + 1)
+
+
+def _read_uint32_be(buffer: bytes, offset: int) -> int:
+    return (
+        _byte(buffer, offset) * 0x1000000
+        + (_byte(buffer, offset + 1) << 16)
+        + (_byte(buffer, offset + 2) << 8)
+        + _byte(buffer, offset + 3)
+    )
+
+
+def _read_uint32_le(buffer: bytes, offset: int) -> int:
+    return (
+        _byte(buffer, offset)
+        + (_byte(buffer, offset + 1) << 8)
+        + (_byte(buffer, offset + 2) << 16)
+        + _byte(buffer, offset + 3) * 0x1000000
+    )
+
+
+def _starts_with(buffer: bytes, values: tuple[int, ...]) -> bool:
+    if len(buffer) < len(values):
+        return False
+    return all(buffer[index] == value for index, value in enumerate(values))
+
+
+def _starts_with_ascii(buffer: bytes, offset: int, text: str) -> bool:
+    if len(buffer) < offset + len(text):
+        return False
+    return all(buffer[offset + index] == ord(character) for index, character in enumerate(text))
+
+
+@dataclass(slots=True)
+class ReaderByteSource:
+    """A :class:`ByteSource` over an environment's ``BinaryReader``."""
+
+    reader: BinaryReader
+    size: int
+    context: Context
+
+    async def read(self, offset: int, length: int) -> bytes:
+        return get_or_throw(await self.reader.read(offset, length, self.context))
+
+
+__all__ = [
+    "BLOCK_BYTES",
+    "HEADER_BYTES",
+    "PNG_SIGNATURE",
+    "ByteSource",
+    "ReaderByteSource",
+    "detect_supported_image_mime_type",
+    "detect_supported_image_mime_type_of",
+    "image_dimensions",
+]

pi_durable/tools/image_processor.py

@@ -0,0 +1,77 @@
+"""Image preparation protocol for the ``read`` tool.
+
+Port of the v1.3.0 ``src/tools/image-processor.ts``. There is no Python image implementation; an
+application supplies an :class:`ImageProcessor` when it wants ``read`` to resize or convert images.
+"""
+
+from __future__ import annotations
+
+import base64
+from typing import Any, Protocol, TypedDict
+
+
+class ImageLimits(TypedDict):
+    """Limits an image the model sees must keep."""
+
+    max_width: int
+    """Pixels. Default 2000."""
+
+    max_height: int
+    """Pixels. Default 2000."""
+
+    max_bytes: int
+    """Bytes of base64. Default 4.5 MiB, below Anthropic's 5 MB."""
+
+
+DEFAULT_IMAGE_LIMITS: ImageLimits = {
+    "max_width": 2000,
+    "max_height": 2000,
+    "max_bytes": int(4.5 * 1024 * 1024),
+}
+
+
+class PreparedImage(TypedDict, total=False):
+    """An image the model can take: a supported format, within the limits, as base64."""
+
+    data: str
+    mime_type: str
+    """The image's format, such as ``image/png``."""
+
+    resized: dict[str, Any]
+    """Set when the image was resized: its size before and after, for the model to map coordinates back."""
+
+    converted_from: str
+    """The original format, set when it differs from ``mime_type``."""
+
+
+class ImageProcessor(Protocol):
+    """Decodes, orients, resizes, and re-encodes images so the model can take them."""
+
+    async def prepare(self, data: bytes, mime_type: str, limits: ImageLimits) -> PreparedImage | None: ...
+
+
+#: Formats every provider takes inline; others, such as BMP, need converting.
+INLINE_IMAGE_TYPES: frozenset[str] = frozenset({"image/png", "image/jpeg", "image/gif", "image/webp"})
+
+
+def to_base64(data: bytes) -> str:
+    """Standard base64 of ``data``."""
+
+    return base64.b64encode(data).decode("ascii")
+
+
+def base64_length(byte_length: int) -> int:
+    """The length of the base64 of ``byte_length`` bytes."""
+
+    return ((byte_length + 2) // 3) * 4
+
+
+__all__ = [
+    "DEFAULT_IMAGE_LIMITS",
+    "INLINE_IMAGE_TYPES",
+    "ImageLimits",
+    "ImageProcessor",
+    "PreparedImage",
+    "base64_length",
+    "to_base64",
+]

pi_durable/tools/path_utils.py

@@ -0,0 +1,111 @@
+"""Tool path normalization and confinement to the execution environment's working directory."""
+
+from __future__ import annotations
+
+import os
+import re
+from collections.abc import Sequence
+
+from ..context import Context
+from ..env import ExecutionEnv, FileError, get_or_throw
+
+UNICODE_SPACES = re.compile("[\u00a0\u2000-\u200a\u202f\u205f\u3000]")
+NARROW_NO_BREAK_SPACE = "\u202f"
+_AM_PM = re.compile(" (AM|PM)\\.", re.IGNORECASE)
+
+
+def normalize_tool_path(path: str) -> str:
+    """Normalize Unicode spaces and drop a leading ``@`` a model may add."""
+
+    normalized = UNICODE_SPACES.sub(" ", path)
+    return normalized[1:] if normalized.startswith("@") else normalized
+
+
+def _within(root: str, path: str) -> bool:
+    try:
+        return os.path.commonpath([root, path]) == root
+    except ValueError:
+        return False
+
+
+async def canonical_path(env: ExecutionEnv, path: str, context: Context) -> str:
+    """The canonical path, or the canonical parent joined with the name when the file is missing."""
+
+    result = await env.canonical_path(path, context)
+    if result.ok:
+        return result.value
+    if result.error.code == "not_supported":
+        return path
+    if result.error.code != "not_found":
+        raise result.error
+    parent = get_or_throw(await env.join_path([path, ".."], context))
+    if parent == path or not path.startswith(parent):
+        return path
+    separator = 0 if parent.endswith(("/", "\\")) else 1
+    name = path[len(parent) + separator :]
+    return get_or_throw(await env.join_path([await canonical_path(env, parent, context), name], context))
+
+
+async def ensure_within_cwd(env: ExecutionEnv, resolved: str, original: str, context: Context) -> str:
+    """Reject a resolved path that leaves the environment's working directory."""
+
+    root = get_or_throw(await env.absolute_path(".", context))
+    if not _within(os.path.abspath(root), os.path.abspath(resolved)):
+        raise ValueError(f"Path {original} is outside the working directory")
+    try:
+        canonical_root = await canonical_path(env, root, context)
+        canonical_target = await canonical_path(env, resolved, context)
+    except FileError as error:
+        if error.code in ("not_supported", "not_found"):
+            # A filesystem without canonical paths falls back to the lexical check above.
+            return resolved
+        raise
+    if not _within(os.path.abspath(canonical_root), os.path.abspath(canonical_target)):
+        raise ValueError(f"Path {original} is outside the working directory")
+    return resolved
+
+
+async def resolve_tool_path(env: ExecutionEnv, path: str, context: Context) -> str:
+    """Resolve the path against the environment, refusing to escape its working directory."""
+
+    resolved = get_or_throw(await env.absolute_path(normalize_tool_path(path), context))
+    return await ensure_within_cwd(env, resolved, path, context)
+
+
+async def resolve_read_tool_path(env: ExecutionEnv, path: str, context: Context) -> str:
+    """Resolve a ``read`` path, trying the Unicode variants a filesystem may use."""
+
+    resolved = await resolve_tool_path(env, path, context)
+    variants: Sequence[str] = (
+        resolved,
+        _AM_PM.sub(lambda match: f"{NARROW_NO_BREAK_SPACE}{match.group(1)}.", resolved),
+        # NFD, for filesystems that store decomposed names.
+        _normalize_nfd(resolved),
+        resolved.replace("'", "\u2019"),
+        _normalize_nfd(resolved.replace("'", "\u2019")),
+    )
+    seen: set[str] = set()
+    for variant in variants:
+        if variant in seen:
+            continue
+        seen.add(variant)
+        if get_or_throw(await env.exists(variant, context)):
+            return variant
+    return resolved
+
+
+def _normalize_nfd(text: str) -> str:
+    import unicodedata
+
+    return unicodedata.normalize("NFD", text)
+
+
+__all__ = [
+    "NARROW_NO_BREAK_SPACE",
+    "UNICODE_SPACES",
+    "canonical_path",
+    "ensure_within_cwd",
+    "normalize_tool_path",
+    "resolve_read_tool_path",
+    "resolve_tool_path",
+]