Compare commits

...

6 Commits

Author SHA1 Message Date
Will Miao
a9e5ee7e79 fix: follow-up nits for AVIF/JXL brotli support
- Fix JXL container ftyp size check (==20 → >=16) to accept
  wider range of valid JXL files
- Add brotli decompression size limit (2 MB) to prevent OOM
- Add trailing newline to requirements.txt
- Add unit tests for new ISOBMFF/brotli extraction paths:
  JXL/AVIF happy paths, missing brob, corrupt payload,
  non-ISOBMFF fallthrough, write-skip on AVIF/JXL,
  JSON dict/list fields, and oversized decompression
2026-06-16 16:27:56 +08:00
Will Miao
a17b0e9901 Merge pull request #982 from koloved/main
Add AVIF and JXL image support with brotli metadata decompression
2026-06-16 16:24:30 +08:00
s.ivanov
8f23d966bf Update requirements.txt 2026-06-16 07:27:32 +02:00
Will Miao
7a76fc72d0 fix(rate-limit): continue to next provider on CivArchive 429 to prevent bulk refresh from freezing (#983)
When CivArchive returns HTTP 429 with a large retry_after, the bulk
metadata refresh would block for hours because:

1. FallbackMetadataProvider raised RateLimitError instead of continuing
   to the next provider (e.g., SQLite archive was never reached).

2. _RateLimitRetryHelper retried long-rate-limit 429s 3 times — all
   futile since the hourly cap hasn't reset.

3. The batch loop had no awareness of persistent rate-limiting,
   causing 192+ models to each hammer the same rate-limited endpoint.

Changes:
- FallbackMetadataProvider: all 6 methods now continue to next provider
  on RateLimitError instead of raising (model_metadata_provider.py)
- fetch_and_update_model: deleted-model path also continues on
  RateLimitError so sqlite provider gets a chance (metadata_sync_service.py)
- _RateLimitRetryHelper: when retry_after >= 120s, only 1 attempt is
  made — retries are futile for hour-scale rate limits
- BulkMetadataRefreshUseCase: tracks consecutive rate-limit failures
  and aborts early after 3 (bulk_metadata_refresh_use_case.py)

Tests: updated test_fallback_respects_retry_limit for new continue
behavior; added tests for large/small retry_after thresholds.
2026-06-16 13:08:34 +08:00
Will Miao
518a4dd5ee chore: add reasonix.toml and .codegraph/ to .gitignore 2026-06-16 13:05:11 +08:00
s.ivanov
2b6d4e5d8b Add AVIF and JXL image support with brotli metadata decompression 2026-06-15 09:28:49 +02:00
13 changed files with 413 additions and 44 deletions

4
.gitignore vendored
View File

@@ -12,12 +12,14 @@ coverage/
.coverage .coverage
model_cache/ model_cache/
# agent # agent / dev tooling
.opencode/ .opencode/
.claude/ .claude/
.sisyphus/ .sisyphus/
.codex .codex
.omo .omo
reasonix.toml
.codegraph/
# Vue widgets development cache (but keep build output) # Vue widgets development cache (but keep build output)
vue-widgets/node_modules/ vue-widgets/node_modules/

View File

@@ -16,6 +16,8 @@ IMG_EXTENSIONS = (
".tif", ".tif",
".tiff", ".tiff",
".webp", ".webp",
".avif",
".jxl",
".mp4" ".mp4"
) )

View File

@@ -216,13 +216,19 @@ class MetadataSyncService:
provider_used: Optional[str] = None provider_used: Optional[str] = None
last_error: Optional[str] = None last_error: Optional[str] = None
civitai_api_not_found = False civitai_api_not_found = False
any_rate_limited = False
for provider_name, provider in provider_attempts: for provider_name, provider in provider_attempts:
try: try:
civitai_metadata_candidate, error = await provider.get_model_by_hash(sha256) civitai_metadata_candidate, error = await provider.get_model_by_hash(sha256)
except RateLimitError as exc: except RateLimitError as exc:
exc.provider = exc.provider or (provider_name or provider.__class__.__name__) logger.warning(
raise "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
provider_name or provider.__class__.__name__,
exc.retry_after or 0,
)
any_rate_limited = True
continue
except Exception as exc: # pragma: no cover - defensive logging except Exception as exc: # pragma: no cover - defensive logging
logger.error("Provider %s failed for hash %s: %s", provider_name, sha256, exc) logger.error("Provider %s failed for hash %s: %s", provider_name, sha256, exc)
civitai_metadata_candidate, error = None, str(exc) civitai_metadata_candidate, error = None, str(exc)
@@ -276,6 +282,8 @@ class MetadataSyncService:
) )
resolved_error = last_error or default_error resolved_error = last_error or default_error
if any_rate_limited and "Rate limited" not in resolved_error:
resolved_error = "Rate limited"
if is_expected_offline_error(resolved_error): if is_expected_offline_error(resolved_error):
resolved_error = OFFLINE_FRIENDLY_MESSAGE resolved_error = OFFLINE_FRIENDLY_MESSAGE

View File

@@ -65,7 +65,14 @@ class _RateLimitRetryHelper:
return await func(*args, **kwargs) return await func(*args, **kwargs)
except RateLimitError as exc: except RateLimitError as exc:
attempt += 1 attempt += 1
if attempt >= self._retry_limit:
# Determine effective retry limit based on rate-limit magnitude
effective_retry_limit = self._retry_limit # default: 3
if exc.retry_after is not None and exc.retry_after >= 120.0:
# Long rate-limit window (>=2 min) — retries are futile
effective_retry_limit = 1 # total 1 attempt = 0 retries
if attempt >= effective_retry_limit:
exc.provider = exc.provider or label exc.provider = exc.provider or label
raise raise
@@ -478,8 +485,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result: if result:
return result, error return result, error
except RateLimitError as exc: except RateLimitError as exc:
exc.provider = exc.provider or label logger.warning(
raise exc "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e: except Exception as e:
logger.debug("Provider %s failed for get_model_by_hash: %s", label, e) logger.debug("Provider %s failed for get_model_by_hash: %s", label, e)
continue continue
@@ -497,16 +508,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result: if result:
return result return result
except RateLimitError as exc: except RateLimitError as exc:
if not_found_confirmed: logger.warning(
logger.debug( "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
"Suppressing rate limit from %s for model %s: " label,
"already confirmed as not found by another provider", exc.retry_after or 0,
label, )
model_id, continue
)
return None
exc.provider = exc.provider or label
raise exc
except ResourceNotFoundError: except ResourceNotFoundError:
not_found_confirmed = True not_found_confirmed = True
logger.debug( logger.debug(
@@ -532,8 +539,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result: if result:
return result return result
except RateLimitError as exc: except RateLimitError as exc:
exc.provider = exc.provider or label logger.warning(
raise exc "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e: except Exception as e:
logger.debug("Provider %s failed for get_model_version: %s", label, e) logger.debug("Provider %s failed for get_model_version: %s", label, e)
continue continue
@@ -550,8 +561,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result: if result:
return result, error return result, error
except RateLimitError as exc: except RateLimitError as exc:
exc.provider = exc.provider or label logger.warning(
raise exc "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e: except Exception as e:
logger.debug("Provider %s failed for get_model_version_info: %s", label, e) logger.debug("Provider %s failed for get_model_version_info: %s", label, e)
continue continue
@@ -572,8 +587,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
except NotImplementedError: except NotImplementedError:
continue continue
except RateLimitError as exc: except RateLimitError as exc:
exc.provider = exc.provider or label logger.warning(
raise exc "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e: except Exception as e:
logger.debug( logger.debug(
"Provider %s failed for get_model_versions_by_hashes: %s", "Provider %s failed for get_model_versions_by_hashes: %s",
@@ -594,8 +613,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result is not None: if result is not None:
return result return result
except RateLimitError as exc: except RateLimitError as exc:
exc.provider = exc.provider or label logger.warning(
raise exc "Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e: except Exception as e:
logger.debug("Provider %s failed for get_user_models: %s", label, e) logger.debug("Provider %s failed for get_user_models: %s", label, e)
continue continue

View File

@@ -77,6 +77,9 @@ class BulkMetadataRefreshUseCase:
await emit("started") await emit("started")
RATE_LIMIT_ABORT_THRESHOLD = 3
consecutive_rate_limits = 0
for model in to_process: for model in to_process:
if self._service.scanner.is_cancelled(): if self._service.scanner.is_cancelled():
self._logger.info("Bulk metadata refresh cancelled by user") self._logger.info("Bulk metadata refresh cancelled by user")
@@ -115,12 +118,39 @@ class BulkMetadataRefreshUseCase:
continue continue
await MetadataManager.hydrate_model_data(model) await MetadataManager.hydrate_model_data(model)
result, _ = await self._metadata_sync.fetch_and_update_model( result, error_msg = await self._metadata_sync.fetch_and_update_model(
sha256=model["sha256"], sha256=model["sha256"],
file_path=model["file_path"], file_path=model["file_path"],
model_data=model, model_data=model,
update_cache_func=self._service.scanner.update_single_model_cache, update_cache_func=self._service.scanner.update_single_model_cache,
) )
if not result and error_msg and "Rate limited" in error_msg:
consecutive_rate_limits += 1
else:
consecutive_rate_limits = 0
if consecutive_rate_limits >= RATE_LIMIT_ABORT_THRESHOLD:
self._logger.warning(
"Bulk metadata refresh aborted: %d consecutive rate limits detected. "
"Processed %d/%d models.",
consecutive_rate_limits,
processed,
total_to_process,
)
await emit(
"rate_limited",
processed=processed,
success=success,
)
return {
"success": False,
"message": f"Rate limit detected; {total_to_process - processed} models skipped",
"processed": processed,
"updated": success,
"total": total_models,
}
if result: if result:
success += 1 success += 1
if original_name != model.get("model_name"): if original_name != model.get("model_name"):

View File

@@ -31,6 +31,8 @@ PREVIEW_EXTENSIONS = [
".mp4", ".mp4",
".gif", ".gif",
".webm", ".webm",
".avif",
".jxl",
] ]
# Card preview image width # Card preview image width
@@ -41,7 +43,7 @@ EXAMPLE_IMAGE_WIDTH = 832
# Supported media extensions for example downloads # Supported media extensions for example downloads
SUPPORTED_MEDIA_EXTENSIONS = { SUPPORTED_MEDIA_EXTENSIONS = {
"images": [".jpg", ".jpeg", ".png", ".webp", ".gif"], "images": [".jpg", ".jpeg", ".png", ".webp", ".gif", ".avif", ".jxl"],
"videos": [".mp4", ".webm"], "videos": [".mp4", ".webm"],
} }

View File

@@ -62,6 +62,10 @@ class ExampleImagesProcessor:
return '.gif' return '.gif'
elif content.startswith(b'RIFF') and b'WEBP' in content[:12]: elif content.startswith(b'RIFF') and b'WEBP' in content[:12]:
return '.webp' return '.webp'
elif len(content) >= 12 and content[4:8] == b'ftyp' and b'avif' in content[8:24]:
return '.avif'
elif content.startswith(b'\x00\x00\x00\x0cJXL \x0d\x0a\x87\x0a'):
return '.jxl'
elif content.startswith(b'\x00\x00\x00\x18ftypmp4') or content.startswith(b'\x00\x00\x00\x20ftypmp4'): elif content.startswith(b'\x00\x00\x00\x18ftypmp4') or content.startswith(b'\x00\x00\x00\x20ftypmp4'):
return '.mp4' return '.mp4'
elif content.startswith(b'\x1A\x45\xDF\xA3'): elif content.startswith(b'\x1A\x45\xDF\xA3'):
@@ -75,6 +79,8 @@ class ExampleImagesProcessor:
'image/png': '.png', 'image/png': '.png',
'image/gif': '.gif', 'image/gif': '.gif',
'image/webp': '.webp', 'image/webp': '.webp',
'image/avif': '.avif',
'image/jxl': '.jxl',
'video/mp4': '.mp4', 'video/mp4': '.mp4',
'video/webm': '.webm', 'video/webm': '.webm',
'video/quicktime': '.mov' 'video/quicktime': '.mov'

View File

@@ -1,17 +1,125 @@
import json import json
import logging import logging
import os import os
import struct
from io import BytesIO from io import BytesIO
from typing import Any, Optional from typing import Any, Optional
import piexif import piexif
from PIL import Image, PngImagePlugin from PIL import Image, PngImagePlugin
try:
import brotli
_BROTLI_AVAILABLE = True
except ImportError:
brotli = None
_BROTLI_AVAILABLE = False
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class ExifUtils: class ExifUtils:
"""Utility functions for working with EXIF data in images""" """Utility functions for working with EXIF data in images"""
@staticmethod
def _parse_isobmff_boxes(data: bytes, offset: int = 0) -> list[dict]:
boxes = []
while offset + 8 <= len(data):
size = struct.unpack('>I', data[offset:offset + 4])[0]
box_type = data[offset + 4:offset + 8]
if size == 0:
break
if size < 8 or offset + size > len(data):
break
box_data = data[offset + 8:offset + size]
boxes.append({'type': box_type, 'data': box_data, 'size': size})
offset += size
return boxes
@staticmethod
def _is_jxl_container(data: bytes) -> bool:
if len(data) < 32:
return False
return (
struct.unpack('>I', data[:4])[0] == 12
and data[4:8] == b'JXL '
and data[8:12] == bytes([0x0d, 0x0a, 0x87, 0x0a])
and struct.unpack('>I', data[12:16])[0] >= 16
and data[16:20] == b'ftyp'
and data[20:24] == b'jxl '
)
@staticmethod
def _is_avif_container(data: bytes) -> bool:
if len(data) < 16:
return False
for box in ExifUtils._parse_isobmff_boxes(data):
if box['type'] == b'ftyp' and b'avif' in box['data']:
return True
return False
# Max decompressed size for brotli metadata (2 MB)
_BROTLI_MAX_DECOMPRESSED = 2 * 1024 * 1024
@staticmethod
def _extract_isobmff_brotli(image_path: str) -> Optional[dict]:
try:
with open(image_path, 'rb') as f:
data = f.read()
except Exception:
return None
if ExifUtils._is_jxl_container(data):
boxes = ExifUtils._parse_isobmff_boxes(data, offset=12)
elif ExifUtils._is_avif_container(data):
boxes = ExifUtils._parse_isobmff_boxes(data)
else:
return None
brob = None
for box in boxes:
if box['type'] == b'brob':
brob = box
break
if brob is None:
return None
payload = brob['data']
if payload[:4] != b'comf':
return None
compressed = payload[4:]
if _BROTLI_AVAILABLE:
try:
decompressed = brotli.decompress(compressed)
if len(decompressed) > ExifUtils._BROTLI_MAX_DECOMPRESSED:
logger.warning(
"Brotli metadata too large (%d bytes, max %d), ignoring",
len(decompressed),
ExifUtils._BROTLI_MAX_DECOMPRESSED,
)
decompressed = None
except Exception:
decompressed = None
else:
decompressed = None
raw = decompressed if decompressed is not None else compressed
try:
meta = json.loads(raw.decode('utf-8'))
except Exception:
return None
result = {"parameters": None, "prompt": None, "workflow": None, "comment": None}
if isinstance(meta.get("prompt"), (dict, list)):
result["prompt"] = json.dumps(meta["prompt"])
elif isinstance(meta.get("prompt"), str):
result["prompt"] = meta["prompt"]
if isinstance(meta.get("workflow"), (dict, list)):
result["workflow"] = json.dumps(meta["workflow"])
elif isinstance(meta.get("workflow"), str):
result["workflow"] = meta["workflow"]
return result
@staticmethod @staticmethod
def _decode_user_comment(user_comment: Any) -> Optional[str]: def _decode_user_comment(user_comment: Any) -> Optional[str]:
if user_comment is None: if user_comment is None:
@@ -43,6 +151,12 @@ class ExifUtils:
"comment": None, "comment": None,
} }
ext = os.path.splitext(image_path)[1].lower()
if ext in ('.avif', '.jxl'):
brotli_meta = ExifUtils._extract_isobmff_brotli(image_path)
if brotli_meta:
return brotli_meta
with Image.open(image_path) as img: with Image.open(image_path) as img:
info = getattr(img, "info", {}) or {} info = getattr(img, "info", {}) or {}
@@ -149,7 +263,6 @@ class ExifUtils:
Optional[str]: Extracted metadata or None if not found Optional[str]: Extracted metadata or None if not found
""" """
try: try:
# Skip for video files
if image_path: if image_path:
ext = os.path.splitext(image_path)[1].lower() ext = os.path.splitext(image_path)[1].lower()
if ext in ['.mp4', '.webm']: if ext in ['.mp4', '.webm']:
@@ -177,10 +290,9 @@ class ExifUtils:
str: Path to the updated image str: Path to the updated image
""" """
try: try:
# Skip for video files
if image_path: if image_path:
ext = os.path.splitext(image_path)[1].lower() ext = os.path.splitext(image_path)[1].lower()
if ext in ['.mp4', '.webm']: if ext in ['.mp4', '.webm', '.avif', '.jxl']:
return image_path return image_path
metadata_fields = ExifUtils._load_structured_metadata(image_path) metadata_fields = ExifUtils._load_structured_metadata(image_path)
@@ -212,10 +324,9 @@ class ExifUtils:
def append_recipe_metadata(image_path, recipe_data) -> str: def append_recipe_metadata(image_path, recipe_data) -> str:
"""Append recipe metadata to an image's EXIF data""" """Append recipe metadata to an image's EXIF data"""
try: try:
# Skip for video files
if image_path: if image_path:
ext = os.path.splitext(image_path)[1].lower() ext = os.path.splitext(image_path)[1].lower()
if ext in ['.mp4', '.webm']: if ext in ['.mp4', '.webm', '.avif', '.jxl']:
return image_path return image_path
# First, extract existing metadata # First, extract existing metadata
@@ -327,10 +438,9 @@ class ExifUtils:
Tuple of (optimized_image_data, extension) Tuple of (optimized_image_data, extension)
""" """
try: try:
# Skip for video files early if it's a file path
if isinstance(image_data, str) and os.path.exists(image_data): if isinstance(image_data, str) and os.path.exists(image_data):
ext = os.path.splitext(image_data)[1].lower() ext = os.path.splitext(image_data)[1].lower()
if ext in ['.mp4', '.webm']: if ext in ['.mp4', '.webm', '.avif', '.jxl']:
try: try:
with open(image_data, 'rb') as f: with open(image_data, 'rb') as f:
return f.read(), ext return f.read(), ext

View File

@@ -13,3 +13,5 @@ aiosqlite
beautifulsoup4 beautifulsoup4
platformdirs platformdirs
pyyaml pyyaml
# brotli — ISOBMFF (AVIF/JXL) metadata decompression
brotli>=1.2.0

View File

@@ -355,9 +355,9 @@ function renderImportInterface(isEmpty) {
<button class="select-files-btn" id="selectExampleFilesBtn"> <button class="select-files-btn" id="selectExampleFilesBtn">
<i class="fas fa-folder-open"></i> Select Files <i class="fas fa-folder-open"></i> Select Files
</button> </button>
<p class="import-formats">Supported formats: jpg, png, gif, webp, mp4, webm</p> <p class="import-formats">Supported formats: jpg, png, gif, webp, avif, jxl, mp4, webm</p>
</div> </div>
<input type="file" id="exampleFilesInput" multiple accept="image/*,video/mp4,video/webm" style="display: none;"> <input type="file" id="exampleFilesInput" multiple accept="image/*,image/avif,image/jxl,video/mp4,video/webm" style="display: none;">
<div class="import-progress-container" style="display: none;"> <div class="import-progress-container" style="display: none;">
<div class="import-progress"> <div class="import-progress">
<div class="progress-bar"></div> <div class="progress-bar"></div>
@@ -473,7 +473,7 @@ export function initExampleImport(modelHash, container) {
*/ */
async function handleImportFiles(files, modelHash, importContainer) { async function handleImportFiles(files, modelHash, importContainer) {
// Filter for supported file types // Filter for supported file types
const supportedImages = ['.jpg', '.jpeg', '.png', '.gif', '.webp']; const supportedImages = ['.jpg', '.jpeg', '.png', '.gif', '.webp', '.avif', '.jxl'];
const supportedVideos = ['.mp4', '.webm']; const supportedVideos = ['.mp4', '.webm'];
const supportedExtensions = [...supportedImages, ...supportedVideos]; const supportedExtensions = [...supportedImages, ...supportedVideos];

View File

@@ -441,7 +441,6 @@ async def test_fetch_and_update_model_returns_rate_limit_error(tmp_path):
assert ok is False assert ok is False
assert error is not None and "Rate limited" in error assert error is not None and "Rate limited" in error
assert "7" in error
helpers.metadata_manager.save_metadata.assert_not_awaited() helpers.metadata_manager.save_metadata.assert_not_awaited()
update_cache.assert_not_awaited() update_cache.assert_not_awaited()
helpers.provider_selector.assert_not_awaited() helpers.provider_selector.assert_not_awaited()

View File

@@ -63,7 +63,8 @@ async def test_fallback_retries_same_provider_on_rate_limit(monkeypatch):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_fallback_respects_retry_limit(monkeypatch): async def test_fallback_continues_to_next_provider_on_rate_limit(monkeypatch):
"""After exhausting retries on primary, fallback should continue to secondary."""
sleep_mock = AsyncMock() sleep_mock = AsyncMock()
monkeypatch.setattr(provider_module.asyncio, "sleep", sleep_mock) monkeypatch.setattr(provider_module.asyncio, "sleep", sleep_mock)
monkeypatch.setattr(provider_module.random, "uniform", lambda *_: 0.0) monkeypatch.setattr(provider_module.random, "uniform", lambda *_: 0.0)
@@ -76,13 +77,13 @@ async def test_fallback_respects_retry_limit(monkeypatch):
rate_limit_retry_limit=2, rate_limit_retry_limit=2,
) )
with pytest.raises(RateLimitError) as exc_info: # After Change A: no longer raises; falls through to secondary
await fallback.get_model_by_hash("abc") result, error = await fallback.get_model_by_hash("abc")
assert exc_info.value.provider == "primary" assert error is None
assert primary.calls == 2 assert result == {"id": "secondary"}
assert secondary.calls == 0 assert primary.calls == 2 # retry_limit exhausted on primary
sleep_mock.assert_awaited_once() assert secondary.calls == 1 # secondary IS called now
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -117,3 +118,40 @@ async def test_rate_limit_retrying_provider_respects_limit(monkeypatch):
assert exc_info.value.provider == "inner" assert exc_info.value.provider == "inner"
assert inner.calls == 2 assert inner.calls == 2
sleep_mock.assert_awaited_once() sleep_mock.assert_awaited_once()
@pytest.mark.asyncio
async def test_retry_helper_limits_retries_for_large_retry_after():
"""With retry_after >= 120s, _RateLimitRetryHelper should only attempt once (no retries)."""
calls = 0
async def failing():
nonlocal calls
calls += 1
raise RateLimitError("limited", retry_after=1500.0)
helper = provider_module._RateLimitRetryHelper(retry_limit=3)
with pytest.raises(RateLimitError):
await helper.run("test", failing)
assert calls == 1 # No retries for large retry_after
@pytest.mark.asyncio
async def test_retry_helper_retries_normally_for_small_retry_after(monkeypatch):
"""With retry_after < 120s, _RateLimitRetryHelper should retry normally (up to limit)."""
sleep_mock = AsyncMock()
monkeypatch.setattr(provider_module.asyncio, "sleep", sleep_mock)
calls = 0
async def succeeding():
nonlocal calls
calls += 1
if calls == 1:
raise RateLimitError("limited", retry_after=30.0)
return {"ok": True}, None
helper = provider_module._RateLimitRetryHelper(retry_limit=3)
result, _ = await helper.run("test", succeeding)
assert result == {"ok": True}
assert calls == 2 # Retried once (small retry_after)

View File

@@ -141,3 +141,150 @@ def test_update_image_metadata_preserves_png_workflow(tmp_path):
img.info["parameters"] img.info["parameters"]
== 'prompt text\nRecipe metadata: {"title":"recipe"}' == 'prompt text\nRecipe metadata: {"title":"recipe"}'
) )
# --- ISOBMFF / brotli extraction tests ---
import struct
import brotli
def _build_jxl_with_brob(payload_json: dict) -> bytes:
"""Build a minimal JXL container with a brob box containing brotli-compressed JSON."""
# ISOBMFF box 1: JXL signature box (size=12, type='JXL ', signature)
box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
# ISOBMFF box 2: ftyp (size=16, type='ftyp', major='jxl ', minor=0)
box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
# ISOBMFF box 3: brob — payload is b'comf' + brotli(json)
compressed = brotli.compress(json.dumps(payload_json).encode("utf-8"))
brob_payload = b"comf" + compressed
box3 = struct.pack(">I", 8 + len(brob_payload)) + b"brob" + brob_payload
return box1 + box2 + box3
def _build_avif_with_brob(payload_json: dict) -> bytes:
"""Build a minimal AVIF container with a brob box containing brotli-compressed JSON."""
compressed = brotli.compress(json.dumps(payload_json).encode("utf-8"))
brob_payload = b"comf" + compressed
ftyp_box = struct.pack(">I", 20) + b"ftyp" + b"avif" + struct.pack(">I", 0) + b"avif"
brob_box = struct.pack(">I", 8 + len(brob_payload)) + b"brob" + brob_payload
return ftyp_box + brob_box
class TestIsobmffBrotliExtraction:
"""Tests for ISOBMFF brotli metadata extraction in ExifUtils."""
def test_extract_jxl_brotli_happy_path(self, tmp_path):
"""JXL container with valid brob box extracts prompt and workflow."""
payload = {"prompt": "a cute cat", "workflow": {"nodes": [{"id": 1}]}}
data = _build_jxl_with_brob(payload)
path = tmp_path / "test.jxl"
path.write_bytes(data)
result = ExifUtils._load_structured_metadata(str(path))
assert result["prompt"] == "a cute cat"
assert result["workflow"] == '{"nodes": [{"id": 1}]}'
assert result["parameters"] is None
assert result["comment"] is None
def test_extract_avif_brotli_happy_path(self, tmp_path):
"""AVIF container with valid brob box extracts prompt and workflow."""
payload = {"prompt": "landscape", "workflow": {"nodes": []}}
data = _build_avif_with_brob(payload)
path = tmp_path / "test.avif"
path.write_bytes(data)
result = ExifUtils._load_structured_metadata(str(path))
assert result["prompt"] == "landscape"
assert result["workflow"] == '{"nodes": []}'
def test_extract_no_brob_box_returns_none(self, tmp_path):
"""JXL container without a brob box returns None from _extract_isobmff_brotli."""
# Only JXL signature + ftyp, no brob
box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
path = tmp_path / "test.jxl"
path.write_bytes(box1 + box2)
# The low-level extraction should return None (no brob box)
result = ExifUtils._extract_isobmff_brotli(str(path))
assert result is None
def test_extract_corrupt_brob_returns_none(self, tmp_path):
"""Broken brob box payload gracefully returns None."""
box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
# brob with garbage payload that doesn't start with b'comf'
garbage = b"\xff\xff\xff\xff" * 32
box3 = struct.pack(">I", 8 + len(garbage)) + b"brob" + garbage
path = tmp_path / "test.jxl"
path.write_bytes(box1 + box2 + box3)
result = ExifUtils._extract_isobmff_brotli(str(path))
assert result is None
def test_extract_non_isobmff_file_falls_through(self, tmp_path):
"""A regular PNG file is not processed as ISOBMFF and returns PIL metadata."""
png_info = PngImagePlugin.PngInfo()
png_info.add_text("prompt", "from png")
path = tmp_path / "test.png"
Image.new("RGB", (4, 4), color="red").save(path, pnginfo=png_info)
result = ExifUtils._load_structured_metadata(str(path))
assert result["prompt"] == "from png"
def test_extract_skip_on_update_and_optimize(self, tmp_path):
"""AVIF/JXL files are skipped for write operations (update/append/optimize)."""
path = tmp_path / "test.avif"
path.write_bytes(b"fake avif data")
# update_image_metadata should return the path unchanged
result = ExifUtils.update_image_metadata(str(path), "some metadata")
assert result == str(path)
# append_recipe_metadata should also skip
result = ExifUtils.append_recipe_metadata(str(path), {"title": "test"})
assert result == str(path)
# optimize_image should passthrough for AVIF/JXL paths
result_data, ext = ExifUtils.optimize_image(str(path))
assert ext == ".avif"
assert result_data == b"fake avif data"
def test_extract_prompt_as_dict(self, tmp_path):
"""prompt field as dict is JSON-serialized."""
payload = {"prompt": {"text": "hello", "negative": "bad"}}
data = _build_jxl_with_brob(payload)
path = tmp_path / "test.jxl"
path.write_bytes(data)
result = ExifUtils._load_structured_metadata(str(path))
assert json.loads(result["prompt"]) == {"text": "hello", "negative": "bad"}
def test_extract_workflow_as_list(self, tmp_path):
"""workflow field as list is JSON-serialized."""
payload = {"workflow": [{"id": 1}, {"id": 2}]}
data = _build_avif_with_brob(payload)
path = tmp_path / "test.avif"
path.write_bytes(data)
result = ExifUtils._load_structured_metadata(str(path))
assert json.loads(result["workflow"]) == [{"id": 1}, {"id": 2}]
def test_over_decompressed_size_limit(self, tmp_path, monkeypatch):
"""Decompressed data exceeding _BROTLI_MAX_DECOMPRESSED is rejected."""
# Monkey-patch the limit to a small value to avoid large test data
monkeypatch.setattr(ExifUtils, "_BROTLI_MAX_DECOMPRESSED", 100)
large_content = "x" * 200
payload = {"prompt": large_content}
data = _build_jxl_with_brob(payload)
path = tmp_path / "test.jxl"
path.write_bytes(data)
# Direct extraction should return None because decompressed size exceeds limit
result = ExifUtils._extract_isobmff_brotli(str(path))
assert result is None