diff --git a/src/google/adk/live/_media_frames.py b/src/google/adk/live/_media_frames.py new file mode 100644 index 0000000000..501466b30d --- /dev/null +++ b/src/google/adk/live/_media_frames.py @@ -0,0 +1,434 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Packs and unpacks live media frame sequences as a single ZIP artifact. + +Live video and image frames arrive as many small blobs. Writing each one as +its own artifact would mean hundreds of round trips to the artifact backend +and would require a media-specific method on ``BaseArtifactService``, which +most backends (in-memory, filesystem, GCS) cannot implement without breaking +their flat ``{filename}/{version}`` layout. + +Instead, the frames plus a ``metadata.json`` manifest are packed into one +uncompressed ZIP and stored with the ordinary ``save_artifact`` call, so +every backend works unchanged. ``ZIP_STORED`` is deliberate: the frames are +already JPEG/PNG so deflate buys nothing, and an uncompressed archive keeps +the central directory usable for reading one member without inflating the +rest. + +Archive layout:: + + .zip + ├── metadata.json + └── frames/ + ├── frame_0000.jpeg + ├── frame_0001.jpeg + └── ... + +Inside ``metadata.json``, each entry in ``frames`` records its archive +``member_name`` under the camelCase JSON key ``fileName`` (for example, +``"frames/frame_0000.jpeg"``). +""" + +from __future__ import annotations + +import io +import json +import logging +import math +import mimetypes +from typing import Any +import zipfile + +from google.genai import types +from pydantic import BaseModel +from pydantic import ConfigDict + +logger = logging.getLogger('google_adk.' + __name__) + +DEFAULT_MIME_TYPE = 'image/jpeg' +"""MIME type assumed when a frame blob does not declare one.""" + +METADATA_FILENAME = 'metadata.json' +"""Name of the manifest member at the root of the archive.""" + +FRAMES_DIRECTORY = 'frames' +"""Name of the archive directory holding the frame members.""" + +MEDIA_ZIP_MIME_TYPE = 'application/zip' +"""MIME type of the packed archive, as stored in the artifact service.""" + +# Python 3.10's stdlib `mimetypes` tables do not include `image/webp` (added in +# Python 3.11). Register it on a private `MimeTypes` instance so Python 3.10 +# containers without `/etc/mime.types` resolve `.webp` without mutating global +# `mimetypes` state. +_MIME_TYPES = mimetypes.MimeTypes() +_MIME_TYPES.add_type('image/webp', '.webp') + + +class MediaFrame(BaseModel): + """A single media frame and the time it was captured.""" + + model_config = ConfigDict( + arbitrary_types_allowed=True, + extra='forbid', + ) + """The pydantic model config.""" + + blob: types.Blob + """The frame payload: its bytes and its MIME type.""" + + timestamp: float + """Wall-clock seconds at which the frame was received.""" + + +def _extension_for_mime_type(mime_type: str | None) -> str: + """Returns the archive file extension to use for a frame MIME type.""" + raw_mime = (mime_type or DEFAULT_MIME_TYPE).split(';')[0].strip().lower() + guessed = _MIME_TYPES.guess_extension(raw_mime, strict=False) + extension = guessed.lstrip('.') if guessed else '' + return 'jpeg' if not extension or extension in ('jpg', 'jpe') else extension + + +def _frame_member_name(index: int, mime_type: str | None) -> str: + """Returns the archive ``member_name`` for the frame at ``index``.""" + extension = _extension_for_mime_type(mime_type) + return f'{FRAMES_DIRECTORY}/frame_{index:04d}.{extension}' + + +def _read_archive_manifest( + archive: zipfile.ZipFile, + member_names: set[str], +) -> dict[str, Any]: + """Reads and validates ``metadata.json`` from an open archive.""" + if METADATA_FILENAME not in member_names: + raise ValueError(f'Media frame archive is missing {METADATA_FILENAME}.') + + manifest = json.loads(archive.read(METADATA_FILENAME).decode('utf-8')) + if not isinstance(manifest, dict): + raise ValueError(f'{METADATA_FILENAME} must contain a JSON object.') + entries = manifest.get('frames', []) + if not isinstance(entries, list): + raise ValueError(f'{METADATA_FILENAME} frames field must be a list.') + start_ts_ms = manifest.get('startTimestampMs', 0) + if ( + isinstance(start_ts_ms, bool) + or not isinstance(start_ts_ms, (int, float)) + or not math.isfinite(start_ts_ms) + ): + raise ValueError( + f'{METADATA_FILENAME} startTimestampMs must be a finite number.' + ) + return manifest + + +def _read_frame_entry( + archive: zipfile.ZipFile, + member_names: set[str], + entry: Any, + start_timestamp: float, +) -> MediaFrame: + """Validates a manifest frame entry and constructs its :class:`MediaFrame`.""" + if not isinstance(entry, dict): + raise ValueError(f'{METADATA_FILENAME} frame entries must be objects.') + member_name = entry.get('fileName') + if ( + not isinstance(member_name, str) + or not member_name + or member_name not in member_names + ): + raise ValueError( + f'Media frame archive is missing frame member {member_name!r}.' + ) + offset_ms = entry.get('offsetMs', 0) + if ( + isinstance(offset_ms, bool) + or not isinstance(offset_ms, (int, float)) + or not math.isfinite(offset_ms) + ): + raise ValueError( + f'{METADATA_FILENAME} frame offsetMs must be a finite number.' + ) + mime_type = entry.get('mimeType', DEFAULT_MIME_TYPE) + if not isinstance(mime_type, str) or not mime_type: + mime_type = DEFAULT_MIME_TYPE + return MediaFrame( + blob=types.Blob( + data=archive.read(member_name), + mime_type=mime_type, + ), + timestamp=start_timestamp + offset_ms / 1000.0, + ) + + +def build_manifest( + frames: list[MediaFrame], + custom_metadata: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Builds the manifest describing a frame sequence. + + The full manifest (including the per-frame index) is written into the + archive as ``metadata.json``, while :func:`summarize_manifest` strips the + ``frames`` index so the bounded summary can be passed to ``save_artifact`` + as ``custom_metadata``. + + Args: + frames: The frames, in capture order. + custom_metadata: Extra caller-supplied keys merged into the manifest. + Computed structural fields (such as ``startTimestampMs``, ``frameCount``, + and ``frames``) always take precedence so the archive remains + self-consistent and unpackable. + + Returns: + The manifest dictionary. + + Raises: + ValueError: If ``frames`` is empty or timestamps are non-finite or + decreasing. + """ + if not frames: + raise ValueError('Cannot build a manifest for an empty frame list.') + + for index, frame in enumerate(frames): + if not math.isfinite(frame.timestamp): + raise ValueError(f'Frame {index} timestamp must be finite.') + + previous_timestamp = frames[0].timestamp + for index, frame in enumerate(frames[1:], start=1): + if frame.timestamp < previous_timestamp: + raise ValueError( + f'Frame timestamps must be non-decreasing (frame {index} at' + f' {frame.timestamp} < {previous_timestamp}).' + ) + previous_timestamp = frame.timestamp + + start_timestamp = frames[0].timestamp + end_timestamp = frames[-1].timestamp + elapsed_seconds = end_timestamp - start_timestamp + # Rounded, not truncated. Seconds-to-milliseconds lands on values like + # 599.9999999999999 for a timestamp that is exactly 0.6s after the start, + # and `int()` would floor that to 599, making every offset up to a + # millisecond short and the sequence fractionally shorter than it was. + duration_ms = round(elapsed_seconds * 1000) + frame_count = len(frames) + # Rate over the gaps between frames, not the frames themselves, so a + # single frame reports 0.0 rather than dividing by a zero interval. + estimated_fps = ( + round((frame_count - 1) / elapsed_seconds, 2) + if frame_count > 1 and elapsed_seconds > 0 + else 0.0 + ) + + frame_entries: list[dict[str, Any]] = [] + for index, frame in enumerate(frames): + data = frame.blob.data or b'' + mime_type = (frame.blob.mime_type or DEFAULT_MIME_TYPE).lower() + frame_entries.append({ + 'frameIndex': index, + 'offsetMs': round((frame.timestamp - start_timestamp) * 1000), + 'fileName': _frame_member_name(index, mime_type), + 'mimeType': mime_type, + 'sizeBytes': len(data), + }) + + manifest: dict[str, Any] = { + **(custom_metadata or {}), + 'type': 'video_frame_sequence', + 'frameCount': frame_count, + 'startTimestampMs': round(start_timestamp * 1000), + 'endTimestampMs': round(end_timestamp * 1000), + 'durationMs': duration_ms, + 'estimatedFps': estimated_fps, + 'frames': frame_entries, + } + return manifest + + +def summarize_manifest(manifest: dict[str, Any]) -> dict[str, Any]: + """Returns the manifest without its per-frame index. + + The per-frame index grows with the frame count and belongs in the archive, + not in the artifact's ``custom_metadata``, which several backends keep in + object metadata with a size limit. The summary is what a caller needs to + decide whether to download the archive at all. + + Args: + manifest: A manifest as returned by :func:`build_manifest`. + + Returns: + A copy of the manifest with the ``frames`` key removed. + """ + return {key: value for key, value in manifest.items() if key != 'frames'} + + +def pack_media_frames( + frames: list[MediaFrame], + custom_metadata: dict[str, Any] | None = None, +) -> tuple[bytes, dict[str, Any]]: + """Packs frames and their manifest into one uncompressed ZIP archive. + + Args: + frames: The frames, in capture order. Must not be empty. + custom_metadata: Extra keys merged into the manifest. + + Returns: + A tuple of the archive bytes and the manifest that was written into it. + + Raises: + ValueError: If ``frames`` is empty or a frame carries no bytes. + """ + if not frames: + raise ValueError('Cannot pack an empty frame list.') + + for index, frame in enumerate(frames): + if not isinstance(frame.blob.data, bytes) or not frame.blob.data: + raise ValueError(f'Frame {index} must contain non-empty byte data.') + + manifest = build_manifest(frames, custom_metadata) + + buffer = io.BytesIO() + # ZIP_STORED: the frames are already compressed image data, so deflating + # them costs CPU for no size win, and storing them keeps each member + # individually readable from the central directory. + with zipfile.ZipFile( + buffer, mode='w', compression=zipfile.ZIP_STORED + ) as archive: + for index, frame in enumerate(frames): + archive.writestr( + _frame_member_name(index, frame.blob.mime_type), + frame.blob.data or b'', + ) + archive.writestr(METADATA_FILENAME, json.dumps(manifest, indent=2)) + + archive_bytes = buffer.getvalue() + logger.debug( + 'Packed %d media frames into a %d byte archive', + len(frames), + len(archive_bytes), + ) + return archive_bytes, manifest + + +def unpack_media_frames( + archive_bytes: bytes, +) -> tuple[list[MediaFrame], dict[str, Any]]: + """Restores the frames and manifest from a packed archive. + + The inverse of :func:`pack_media_frames`: every frame comes back with the + bytes and MIME type it went in with. Capture times are stored as + whole-millisecond offsets from the first frame, so they come back rounded + to the nearest millisecond rather than bit-for-bit. + + Args: + archive_bytes: The archive body, as returned by ``load_artifact``. + + Returns: + A tuple of the frames in capture order and the manifest. + + Raises: + ValueError: If the archive is not a valid ZIP, has a corrupt member CRC, + has no manifest, or names a frame member that is not present. + """ + try: + with zipfile.ZipFile(io.BytesIO(archive_bytes), mode='r') as archive: + member_names = set(archive.namelist()) + manifest = _read_archive_manifest(archive, member_names) + start_timestamp = manifest.get('startTimestampMs', 0) / 1000.0 + frames = [ + _read_frame_entry(archive, member_names, entry, start_timestamp) + for entry in manifest.get('frames', []) + ] + except ( + zipfile.BadZipFile, + json.JSONDecodeError, + UnicodeDecodeError, + ) as exc: + raise ValueError('Media frame archive is not a valid ZIP.') from exc + + return frames, manifest + + +def read_manifest(archive_bytes: bytes) -> dict[str, Any]: + """Reads only the manifest out of a packed archive. + + Args: + archive_bytes: The archive body, as returned by ``load_artifact``. + + Returns: + The manifest dictionary. + + Raises: + ValueError: If the archive is not a valid ZIP, has a corrupt member CRC, + or has no manifest. + """ + try: + with zipfile.ZipFile(io.BytesIO(archive_bytes), mode='r') as archive: + return _read_archive_manifest(archive, set(archive.namelist())) + except ( + zipfile.BadZipFile, + json.JSONDecodeError, + UnicodeDecodeError, + ) as exc: + raise ValueError('Media frame archive is not a valid ZIP.') from exc + + +def extract_frame(archive_bytes: bytes, index: int) -> MediaFrame: + """Reads a single frame out of a packed archive. + + Only the requested member is read; the rest of the archive is located + through the ZIP central directory and never decoded. + + Args: + archive_bytes: The archive body, as returned by ``load_artifact``. + index: Zero-based position of the frame in capture order. + + Returns: + The requested frame. + + Raises: + ValueError: If the archive is malformed, has no manifest, or has no frame + at ``index``. + """ + try: + with zipfile.ZipFile(io.BytesIO(archive_bytes), mode='r') as archive: + member_names = set(archive.namelist()) + manifest = _read_archive_manifest(archive, member_names) + entries = manifest.get('frames', []) + if index < 0 or index >= len(entries): + raise ValueError(f'Media frame archive has no frame at index {index}.') + start_timestamp = manifest.get('startTimestampMs', 0) / 1000.0 + return _read_frame_entry( + archive, member_names, entries[index], start_timestamp + ) + except ( + zipfile.BadZipFile, + json.JSONDecodeError, + UnicodeDecodeError, + ) as exc: + raise ValueError('Media frame archive is not a valid ZIP.') from exc + + +def extract_preview_frame(archive_bytes: bytes) -> MediaFrame: + """Reads the first frame out of a packed archive. + + Args: + archive_bytes: The archive body, as returned by ``load_artifact``. + + Returns: + The first frame in capture order. + + Raises: + ValueError: If the archive is malformed or holds no frames. + """ + return extract_frame(archive_bytes, 0) diff --git a/tests/unittests/live/test_media_frames.py b/tests/unittests/live/test_media_frames.py new file mode 100644 index 0000000000..e55cb54f40 --- /dev/null +++ b/tests/unittests/live/test_media_frames.py @@ -0,0 +1,751 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for packing live media frames into a single ZIP artifact.""" + +from __future__ import annotations + +import io +import json +import zipfile + +from google.adk.live._media_frames import build_manifest +from google.adk.live._media_frames import DEFAULT_MIME_TYPE +from google.adk.live._media_frames import extract_frame +from google.adk.live._media_frames import extract_preview_frame +from google.adk.live._media_frames import FRAMES_DIRECTORY +from google.adk.live._media_frames import MEDIA_ZIP_MIME_TYPE +from google.adk.live._media_frames import MediaFrame +from google.adk.live._media_frames import METADATA_FILENAME +from google.adk.live._media_frames import pack_media_frames +from google.adk.live._media_frames import read_manifest +from google.adk.live._media_frames import summarize_manifest +from google.adk.live._media_frames import unpack_media_frames +from google.genai import types +import pydantic +import pytest + + +def make_frames( + count: int = 3, + *, + mime_type: str = 'image/jpeg', + start: float = 100.0, + step: float = 0.2, +) -> list[MediaFrame]: + """Builds a frame batch whose bytes identify their index.""" + return [ + MediaFrame( + blob=types.Blob( + data=f'frame-{index}-payload'.encode(), mime_type=mime_type + ), + timestamp=start + index * step, + ) + for index in range(count) + ] + + +class TestMediaFrame: + """The model is the contract between the cache manager and the packer.""" + + def test_unknown_fields_are_rejected(self) -> None: + # `extra="forbid"` is what stops a typo'd field name from being silently + # dropped on its way into the archive. + with pytest.raises(pydantic.ValidationError): + MediaFrame( + blob=types.Blob(data=b'x', mime_type='image/jpeg'), + timestamp=0.0, + **{'not_a_field': 1}, # type: ignore[arg-type] + ) + + def test_accepts_the_declared_fields(self) -> None: + frame = MediaFrame( + blob=types.Blob(data=b'x', mime_type='image/png'), timestamp=1.5 + ) + assert frame.blob.mime_type == 'image/png' + assert frame.timestamp == 1.5 + + +class TestRoundTrip: + """pack and unpack have to be exact inverses, or stored video is lossy.""" + + def test_unpack_returns_the_frames_that_were_packed(self) -> None: + frames = make_frames(5) + + archive_bytes, _ = pack_media_frames(frames) + restored, _ = unpack_media_frames(archive_bytes) + + assert restored == frames + + def test_round_trip_preserves_mixed_mime_types(self) -> None: + frames = [ + MediaFrame( + blob=types.Blob(data=b'jpeg-bytes', mime_type='image/jpeg'), + timestamp=1.0, + ), + MediaFrame( + blob=types.Blob(data=b'png-bytes', mime_type='image/png'), + timestamp=2.0, + ), + MediaFrame( + blob=types.Blob(data=b'webp-bytes', mime_type='image/webp'), + timestamp=3.0, + ), + ] + + archive_bytes, _ = pack_media_frames(frames) + restored, _ = unpack_media_frames(archive_bytes) + + assert [frame.blob.mime_type for frame in restored] == [ + 'image/jpeg', + 'image/png', + 'image/webp', + ] + assert restored == frames + + def test_round_trip_preserves_a_single_frame(self) -> None: + frames = make_frames(1) + + archive_bytes, _ = pack_media_frames(frames) + restored, _ = unpack_media_frames(archive_bytes) + + assert restored == frames + + def test_round_trip_preserves_frame_order_for_a_long_sequence(self) -> None: + """Members are named with a zero-padded index precisely so that frame 10 + does not sort before frame 2.""" + frames = make_frames(25, step=0.04) + + archive_bytes, _ = pack_media_frames(frames) + restored, _ = unpack_media_frames(archive_bytes) + + assert [frame.blob.data for frame in restored] == [ + frame.blob.data for frame in frames + ] + + def test_timestamps_round_trip_to_millisecond_precision(self) -> None: + """Capture times are stored as whole-millisecond offsets from the start, + so a round trip quantises them. Recording the limit here keeps a caller + from assuming microsecond fidelity it will not get.""" + frames = [ + MediaFrame( + blob=types.Blob(data=b'a', mime_type='image/jpeg'), + timestamp=100.00049, + ), + MediaFrame( + blob=types.Blob(data=b'b', mime_type='image/jpeg'), + timestamp=100.30051, + ), + ] + + archive_bytes, _ = pack_media_frames(frames) + restored, _ = unpack_media_frames(archive_bytes) + + assert restored[0].timestamp == pytest.approx(100.0, abs=1e-3) + assert restored[1].timestamp == pytest.approx(100.3, abs=1e-3) + + +class TestArchiveLayout: + """The on-disk shape is a contract with any reader that is not this module.""" + + def test_archive_holds_the_frames_and_the_manifest(self) -> None: + frames = make_frames(3) + + archive_bytes, _ = pack_media_frames(frames) + + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: + assert sorted(archive.namelist()) == [ + f'{FRAMES_DIRECTORY}/frame_0000.jpeg', + f'{FRAMES_DIRECTORY}/frame_0001.jpeg', + f'{FRAMES_DIRECTORY}/frame_0002.jpeg', + METADATA_FILENAME, + ] + for index, frame in enumerate(frames): + member = f'{FRAMES_DIRECTORY}/frame_{index:04d}.jpeg' + assert archive.read(member) == frame.blob.data + + def test_frames_are_stored_uncompressed(self) -> None: + """Frames arrive already JPEG/PNG encoded, so deflate would spend CPU per + frame for no size win, and storing keeps members individually readable.""" + archive_bytes, _ = pack_media_frames(make_frames(3)) + + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: + for info in archive.infolist(): + assert info.compress_type == zipfile.ZIP_STORED + + def test_manifest_member_is_valid_json(self) -> None: + archive_bytes, manifest = pack_media_frames(make_frames(2)) + + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: + written = json.loads(archive.read(METADATA_FILENAME).decode('utf-8')) + + assert written == manifest + + @pytest.mark.parametrize( + 'mime_type,expected_extension', + [ + ('image/jpeg', 'jpeg'), + # jpg and jpeg are the same encoding; collapsing them keeps one + # sequence from mixing two names for one format. + ('image/jpg', 'jpeg'), + ('image/png', 'png'), + ('image/webp', 'webp'), + ('image/gif', 'gif'), + ('image/svg+xml', 'svg'), + ('video/mp4', 'mp4'), + ('IMAGE/PNG', 'png'), + ('image/png;codecs=foo', 'png'), + ('image/png ', 'png'), + # Anything that is not a recognized MIME type falls back rather than + # failing the batch: an odd mime type is not worth losing video over. + ('notamimetype', 'jpeg'), + ('image/', 'jpeg'), + ('', 'jpeg'), + (None, 'jpeg'), + ], + ) + def test_member_extension_follows_the_frame_mime_type( + self, mime_type: str | None, expected_extension: str + ) -> None: + frames = [ + MediaFrame( + blob=types.Blob(data=b'x', mime_type=mime_type), timestamp=0.0 + ) + ] + + archive_bytes, _ = pack_media_frames(frames) + + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: + assert archive.namelist()[0] == ( + f'{FRAMES_DIRECTORY}/frame_0000.{expected_extension}' + ) + + @pytest.mark.parametrize( + 'malicious_mime', + [ + 'image/../../etc/passwd', + 'image/..\\..\\etc\\passwd', + 'image/..', + ], + ) + def test_a_mime_type_carrying_separators_cannot_escape_the_frames_directory( + self, malicious_mime: str + ) -> None: + """The mime type arrives from a live model stream, so it is untrusted + input that ends up in a member name.""" + frames = [ + MediaFrame( + blob=types.Blob(data=b'x', mime_type=malicious_mime), + timestamp=0.0, + ) + ] + + archive_bytes, _ = pack_media_frames(frames) + + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: + member = archive.namelist()[0] + assert member == f'{FRAMES_DIRECTORY}/frame_0000.jpeg' + assert '..' not in member + assert '\\' not in member + assert member.count('/') == 1 + + +class TestManifest: + """The manifest is what a reader consults before downloading anything.""" + + def test_manifest_describes_the_sequence(self) -> None: + frames = make_frames(5, start=10.0, step=0.25) + + manifest = build_manifest(frames) + + assert manifest['type'] == 'video_frame_sequence' + assert manifest['frameCount'] == 5 + assert manifest['startTimestampMs'] == 10000 + assert manifest['endTimestampMs'] == 11000 + assert manifest['durationMs'] == 1000 + assert manifest['estimatedFps'] == 4.0 + + def test_manifest_describes_every_frame(self) -> None: + frames = make_frames(4, start=10.0, step=0.25) + + manifest = build_manifest(frames) + + assert [entry['frameIndex'] for entry in manifest['frames']] == [0, 1, 2, 3] + assert [entry['offsetMs'] for entry in manifest['frames']] == [ + 0, + 250, + 500, + 750, + ] + assert [entry['fileName'] for entry in manifest['frames']] == [ + f'{FRAMES_DIRECTORY}/frame_{index:04d}.jpeg' for index in range(4) + ] + assert [entry['mimeType'] for entry in manifest['frames']] == [ + 'image/jpeg' + ] * 4 + assert [entry['sizeBytes'] for entry in manifest['frames']] == [ + len(frame.blob.data or b'') for frame in frames + ] + + def test_single_frame_reports_no_frame_rate(self) -> None: + """One frame spans no interval, so a rate would be a division by zero + dressed up as data.""" + manifest = build_manifest(make_frames(1)) + + assert manifest['frameCount'] == 1 + assert manifest['durationMs'] == 0 + assert manifest['estimatedFps'] == 0.0 + + def test_frames_sharing_one_timestamp_report_no_frame_rate(self) -> None: + """Frames captured inside the same clock tick are legitimate and must not + divide by a zero interval.""" + manifest = build_manifest(make_frames(3, step=0.0)) + + assert manifest['durationMs'] == 0 + assert manifest['estimatedFps'] == 0.0 + assert [entry['offsetMs'] for entry in manifest['frames']] == [0, 0, 0] + + def test_frame_rate_counts_intervals_not_frames(self) -> None: + """Five frames over one second is four intervals, so 4fps, not 5.""" + manifest = build_manifest(make_frames(5, start=0.0, step=0.25)) + + assert manifest['estimatedFps'] == 4.0 + + def test_frame_without_a_mime_type_is_recorded_as_the_default(self) -> None: + frames = [ + MediaFrame(blob=types.Blob(data=b'x', mime_type=None), timestamp=0.0) + ] + + manifest = build_manifest(frames) + + assert manifest['frames'][0]['mimeType'] == DEFAULT_MIME_TYPE + + def test_mime_types_are_recorded_lowercased(self) -> None: + frames = [ + MediaFrame( + blob=types.Blob(data=b'x', mime_type='IMAGE/PNG'), timestamp=0.0 + ) + ] + + manifest = build_manifest(frames) + + assert manifest['frames'][0]['mimeType'] == 'image/png' + + def test_empty_frame_list_is_rejected(self) -> None: + with pytest.raises(ValueError, match='empty frame list'): + build_manifest([]) + + def test_decreasing_timestamps_are_rejected(self) -> None: + frames = [ + MediaFrame( + blob=types.Blob(data=b'a', mime_type='image/jpeg'), timestamp=10.0 + ), + MediaFrame( + blob=types.Blob(data=b'b', mime_type='image/jpeg'), timestamp=9.5 + ), + ] + + with pytest.raises(ValueError, match='must be non-decreasing'): + build_manifest(frames) + + def test_custom_metadata_is_merged_in(self) -> None: + manifest = build_manifest( + make_frames(2), custom_metadata={'source': 'webcam'} + ) + + assert manifest['source'] == 'webcam' + assert manifest['frameCount'] == 2 + + def test_custom_metadata_cannot_override_computed_keys(self) -> None: + """Computed structural keys (`startTimestampMs`, `frameCount`, `frames`, + etc.) take precedence over `custom_metadata` so unpacked timestamps and + members cannot be shifted or corrupted.""" + frames = make_frames(2, start=100.0, step=0.2) + + archive_bytes, manifest = pack_media_frames( + frames, + custom_metadata={ + 'source': 'webcam', + 'frameCount': 999, + 'startTimestampMs': 0, + }, + ) + + assert manifest['source'] == 'webcam' + assert manifest['frameCount'] == 2 + assert manifest['startTimestampMs'] == 100000 + restored, _ = unpack_media_frames(archive_bytes) + assert restored == frames + assert extract_frame(archive_bytes, 1) == frames[1] + + def test_overridden_frames_key_cannot_misname_or_break_unpack(self) -> None: + """`custom_metadata` cannot override `frames`, so the archive always + records and unpacks the packed members.""" + frames = make_frames(2) + + archive_bytes, manifest = pack_media_frames( + frames, custom_metadata={'frames': [{'fileName': 'evil/path.jpeg'}]} + ) + + assert [entry['fileName'] for entry in manifest['frames']] == [ + f'{FRAMES_DIRECTORY}/frame_0000.jpeg', + f'{FRAMES_DIRECTORY}/frame_0001.jpeg', + ] + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: + names = archive.namelist() + assert f'{FRAMES_DIRECTORY}/frame_0000.jpeg' in names + assert f'{FRAMES_DIRECTORY}/frame_0001.jpeg' in names + assert 'evil/path.jpeg' not in names + restored, unpacked_manifest = unpack_media_frames(archive_bytes) + assert restored == frames + assert unpacked_manifest == manifest + assert extract_frame(archive_bytes, 0) == frames[0] + + +class TestSummarizeManifest: + """What travels in artifact metadata has to stay small and bounded.""" + + def test_summary_drops_the_per_frame_index(self) -> None: + manifest = build_manifest(make_frames(3)) + + summary = summarize_manifest(manifest) + + assert 'frames' not in summary + assert summary['frameCount'] == 3 + assert summary['type'] == 'video_frame_sequence' + assert summary['durationMs'] == manifest['durationMs'] + assert summary['estimatedFps'] == manifest['estimatedFps'] + + def test_summary_does_not_mutate_the_manifest(self) -> None: + manifest = build_manifest(make_frames(3)) + + summarize_manifest(manifest) + + assert len(manifest['frames']) == 3 + + def test_summary_size_does_not_grow_with_the_frame_count(self) -> None: + """Several artifact backends keep custom_metadata in object metadata with + a few-kilobyte cap, which a per-frame index would blow through within a + minute of video.""" + short_summary = summarize_manifest(build_manifest(make_frames(2))) + long_summary = summarize_manifest(build_manifest(make_frames(600))) + + assert short_summary.keys() == long_summary.keys() + assert len(json.dumps(long_summary).encode()) < 1024 + + +class TestPackValidation: + """Nothing half-formed should reach the archive.""" + + def test_empty_frame_list_is_rejected(self) -> None: + with pytest.raises(ValueError, match='empty frame list'): + pack_media_frames([]) + + def test_decreasing_timestamps_are_rejected(self) -> None: + frames = [ + MediaFrame( + blob=types.Blob(data=b'a', mime_type='image/jpeg'), timestamp=5.0 + ), + MediaFrame( + blob=types.Blob(data=b'b', mime_type='image/jpeg'), timestamp=4.9 + ), + ] + + with pytest.raises(ValueError, match='must be non-decreasing'): + pack_media_frames(frames) + + def test_frame_with_no_bytes_is_rejected(self) -> None: + frames = make_frames(3) + frames[1] = MediaFrame( + blob=types.Blob(data=b'', mime_type='image/jpeg'), timestamp=1.0 + ) + + with pytest.raises(ValueError, match='Frame 1 must contain non-empty'): + pack_media_frames(frames) + + def test_frame_with_none_bytes_is_rejected(self) -> None: + frames = [ + MediaFrame( + blob=types.Blob(data=None, mime_type='image/jpeg'), timestamp=1.0 + ) + ] + + with pytest.raises(ValueError, match='Frame 0 must contain non-empty'): + pack_media_frames(frames) + + +class TestUnpackValidation: + """A corrupt archive must say so rather than return partial video.""" + + def test_non_zip_input_is_rejected(self) -> None: + with pytest.raises(ValueError, match='not a valid ZIP'): + unpack_media_frames(b'this is not a zip file') + + def test_archive_without_a_manifest_is_rejected(self) -> None: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode='w') as archive: + archive.writestr(f'{FRAMES_DIRECTORY}/frame_0000.jpeg', b'x') + + with pytest.raises(ValueError, match=f'missing {METADATA_FILENAME}'): + unpack_media_frames(buffer.getvalue()) + + def test_manifest_naming_an_absent_member_is_rejected(self) -> None: + """Silently skipping the missing frame would hand back a sequence shorter + than its own manifest claims.""" + frames = make_frames(3) + archive_bytes, manifest = pack_media_frames(frames) + + rebuilt = io.BytesIO() + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as source: + with zipfile.ZipFile(rebuilt, mode='w') as target: + for name in source.namelist(): + if name.endswith('frame_0001.jpeg'): + continue + target.writestr(name, source.read(name)) + + with pytest.raises(ValueError, match='missing frame member'): + unpack_media_frames(rebuilt.getvalue()) + + def test_non_dict_manifest_raises_value_error_across_all_readers( + self, + ) -> None: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode='w') as archive: + archive.writestr(f'{FRAMES_DIRECTORY}/frame_0000.jpeg', b'x') + archive.writestr(METADATA_FILENAME, json.dumps([1, 2, 3])) + corrupt_bytes = buffer.getvalue() + + with pytest.raises(ValueError, match='must contain a JSON object'): + unpack_media_frames(corrupt_bytes) + with pytest.raises(ValueError, match='must contain a JSON object'): + read_manifest(corrupt_bytes) + with pytest.raises(ValueError, match='must contain a JSON object'): + extract_frame(corrupt_bytes, 0) + + def test_bad_member_crc_raises_value_error_across_all_readers(self) -> None: + """`ZipFile(...)` only validates the central directory; a corrupt member + payload raises `BadZipFile` later inside `archive.read()`, which must still + surface as the documented `ValueError`.""" + archive_bytes, _ = pack_media_frames(make_frames(2)) + corrupted = bytearray(archive_bytes) + payload_offset = corrupted.index(b'frame-0-payload') + corrupted[payload_offset] ^= 0xFF + corrupt_frame_bytes = bytes(corrupted) + + with pytest.raises(ValueError, match='not a valid ZIP'): + unpack_media_frames(corrupt_frame_bytes) + with pytest.raises(ValueError, match='not a valid ZIP'): + extract_frame(corrupt_frame_bytes, 0) + + corrupted_manifest = bytearray(archive_bytes) + manifest_offset = corrupted_manifest.index(b'video_frame_sequence') + corrupted_manifest[manifest_offset] ^= 0xFF + corrupt_manifest_bytes = bytes(corrupted_manifest) + + with pytest.raises(ValueError, match='not a valid ZIP'): + read_manifest(corrupt_manifest_bytes) + + +class TestReadManifest: + """Reading the shape of a sequence should not cost the whole sequence.""" + + def test_read_manifest_matches_the_packed_manifest(self) -> None: + archive_bytes, manifest = pack_media_frames(make_frames(4)) + + assert read_manifest(archive_bytes) == manifest + + def test_non_zip_input_is_rejected(self) -> None: + with pytest.raises(ValueError, match='not a valid ZIP'): + read_manifest(b'not a zip') + + def test_archive_without_a_manifest_is_rejected(self) -> None: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode='w') as archive: + archive.writestr('unrelated.txt', b'x') + + with pytest.raises(ValueError, match=f'missing {METADATA_FILENAME}'): + read_manifest(buffer.getvalue()) + + +class TestExtractFrame: + """Random access is the reason for ZIP over TAR.""" + + def test_extract_returns_the_requested_frame(self) -> None: + frames = make_frames(5) + archive_bytes, _ = pack_media_frames(frames) + + assert extract_frame(archive_bytes, 2) == frames[2] + + def test_preview_is_the_first_frame(self) -> None: + frames = make_frames(4) + archive_bytes, _ = pack_media_frames(frames) + + assert extract_preview_frame(archive_bytes) == frames[0] + + def test_extract_reads_only_the_requested_member(self) -> None: + """A reader wanting a thumbnail should not have to hold every frame in + memory; the central directory is what makes that possible. An archive + stripped of every other frame still yields frame 0.""" + frames = make_frames(5) + archive_bytes, _ = pack_media_frames(frames) + + trimmed = io.BytesIO() + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as source: + with zipfile.ZipFile(trimmed, mode='w') as target: + target.writestr(METADATA_FILENAME, source.read(METADATA_FILENAME)) + first = f'{FRAMES_DIRECTORY}/frame_0000.jpeg' + target.writestr(first, source.read(first)) + + assert extract_frame(trimmed.getvalue(), 0) == frames[0] + + @pytest.mark.parametrize('index', [-1, 3, 99]) + def test_out_of_range_index_is_rejected(self, index: int) -> None: + archive_bytes, _ = pack_media_frames(make_frames(3)) + + with pytest.raises(ValueError, match='no frame at index'): + extract_frame(archive_bytes, index) + + def test_non_zip_input_is_rejected(self) -> None: + with pytest.raises(ValueError, match='not a valid ZIP'): + extract_frame(b'not a zip', 0) + + def test_manifest_naming_an_absent_member_is_rejected(self) -> None: + frames = make_frames(2) + archive_bytes, _ = pack_media_frames(frames) + + rebuilt = io.BytesIO() + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as source: + with zipfile.ZipFile(rebuilt, mode='w') as target: + target.writestr(METADATA_FILENAME, source.read(METADATA_FILENAME)) + + with pytest.raises(ValueError, match='missing frame member'): + extract_frame(rebuilt.getvalue(), 0) + + def test_archive_without_a_manifest_is_rejected(self) -> None: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode='w') as archive: + archive.writestr(f'{FRAMES_DIRECTORY}/frame_0000.jpeg', b'x') + + with pytest.raises(ValueError, match=f'missing {METADATA_FILENAME}'): + extract_frame(buffer.getvalue(), 0) + + def test_five_digit_frame_index_preserves_manifest_order(self) -> None: + """Wide indices (`frame_10000.jpeg` after `frame_9999.jpeg`) are resolved + directly from the `frames` list in `metadata.json` rather than by + lexicographic member sorting.""" + manifest = { + 'type': 'video_frame_sequence', + 'frameCount': 3, + 'startTimestampMs': 1000, + 'endTimestampMs': 3000, + 'durationMs': 2000, + 'estimatedFps': 1.0, + 'frames': [ + { + 'frameIndex': 0, + 'offsetMs': 0, + 'fileName': f'{FRAMES_DIRECTORY}/frame_0000.jpeg', + 'mimeType': 'image/jpeg', + 'sizeBytes': 2, + }, + { + 'frameIndex': 9999, + 'offsetMs': 1000, + 'fileName': f'{FRAMES_DIRECTORY}/frame_9999.jpeg', + 'mimeType': 'image/jpeg', + 'sizeBytes': 2, + }, + { + 'frameIndex': 10000, + 'offsetMs': 2000, + 'fileName': f'{FRAMES_DIRECTORY}/frame_10000.jpeg', + 'mimeType': 'image/jpeg', + 'sizeBytes': 2, + }, + ], + } + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode='w') as archive: + archive.writestr(f'{FRAMES_DIRECTORY}/frame_10000.jpeg', b'f2') + archive.writestr(f'{FRAMES_DIRECTORY}/frame_9999.jpeg', b'f1') + archive.writestr(f'{FRAMES_DIRECTORY}/frame_0000.jpeg', b'f0') + archive.writestr(METADATA_FILENAME, json.dumps(manifest)) + archive_bytes = buffer.getvalue() + + assert [extract_frame(archive_bytes, i).blob.data for i in range(3)] == [ + b'f0', + b'f1', + b'f2', + ] + + def test_non_frame_files_inside_frames_directory_are_ignored(self) -> None: + frames = make_frames(1) + archive_bytes, _ = pack_media_frames(frames) + + rebuilt = io.BytesIO() + with zipfile.ZipFile(io.BytesIO(archive_bytes)) as source: + with zipfile.ZipFile(rebuilt, mode='w') as target: + target.writestr(f'{FRAMES_DIRECTORY}/notes.txt', b'stray note') + for member_name in source.namelist(): + target.writestr(member_name, source.read(member_name)) + archive_with_stray = rebuilt.getvalue() + + assert extract_frame(archive_with_stray, 0) == frames[0] + with pytest.raises(ValueError, match='no frame at index'): + extract_frame(archive_with_stray, 1) + + @pytest.mark.parametrize( + 'bad_manifest,expected_match', + [ + ({'frames': 'not-a-list'}, 'frames field must be a list'), + ({'frames': ['not-a-dict']}, 'frame entries must be objects'), + ( + { + 'startTimestampMs': 'bad', + 'frames': [{'fileName': 'frames/frame_0000.jpeg'}], + }, + 'startTimestampMs must be a finite number', + ), + ( + { + 'frames': [{ + 'fileName': 'frames/frame_0000.jpeg', + 'offsetMs': 'bad', + }] + }, + 'offsetMs must be a finite number', + ), + ], + ) + def test_shared_manifest_and_entry_validation_across_readers( + self, + bad_manifest: dict[str, object], + expected_match: str, + ) -> None: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode='w') as archive: + archive.writestr(f'{FRAMES_DIRECTORY}/frame_0000.jpeg', b'x') + archive.writestr(METADATA_FILENAME, json.dumps(bad_manifest)) + corrupt_bytes = buffer.getvalue() + + with pytest.raises(ValueError, match=expected_match): + unpack_media_frames(corrupt_bytes) + with pytest.raises(ValueError, match=expected_match): + extract_frame(corrupt_bytes, 0) + + +def test_zip_mime_type_is_the_one_stored_on_the_artifact() -> None: + """The declared type has to describe what `load_artifact` hands back, which + is the archive rather than any frame inside it.""" + assert MEDIA_ZIP_MIME_TYPE == 'application/zip'