From de3d0571f809e9187600b6fd4fcdc3ff61eb34de Mon Sep 17 00:00:00 2001 From: Will Miao Date: Wed, 25 Mar 2026 20:33:43 +0800 Subject: [PATCH] fix: verify returned image ID matches requested ID in CivitAI API Fix issue #870 where importing recipes from CivitAI image URLs would return the wrong image when the API response did not contain the requested image ID. The get_image_info() method now: - Iterates through all returned items to find matching ID - Returns None when no match is found and logs warning with returned IDs - Handles invalid (non-numeric) ID formats New test cases: - test_get_image_info_returns_matching_item - test_get_image_info_returns_none_when_id_mismatch - test_get_image_info_handles_invalid_id Co-Authored-By: Claude Opus 4.6 --- py/services/civitai_client.py | 29 ++++++++++++++++++--- tests/services/test_civitai_client.py | 36 ++++++++++++++++++++++++--- 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/py/services/civitai_client.py b/py/services/civitai_client.py index 31a51aed..2bdd46c7 100644 --- a/py/services/civitai_client.py +++ b/py/services/civitai_client.py @@ -490,14 +490,33 @@ class CivitaiClient: """ try: url = f"{self.base_url}/images?imageId={image_id}&nsfw=X" + requested_id = int(image_id) logger.debug(f"Fetching image info for ID: {image_id}") success, result = await self._make_request("GET", url, use_auth=True) if success: - if result and "items" in result and len(result["items"]) > 0: - logger.debug(f"Successfully fetched image info for ID: {image_id}") - return result["items"][0] + if result and "items" in result and isinstance(result["items"], list): + items = result["items"] + + # First, try to find the item with matching ID + for item in items: + if isinstance(item, dict) and item.get("id") == requested_id: + logger.debug(f"Successfully fetched image info for ID: {image_id}") + return item + + # No matching ID found - log warning with details about returned items + returned_ids = [ + item.get("id") for item in items + if isinstance(item, dict) and "id" in item + ] + logger.warning( + f"CivitAI API returned no matching image for requested ID {image_id}. " + f"Returned {len(items)} item(s) with IDs: {returned_ids}. " + f"This may indicate the image was deleted, hidden, or there is a database lag." + ) + return None + logger.warning(f"No image found with ID: {image_id}") return None @@ -505,6 +524,10 @@ class CivitaiClient: return None except RateLimitError: raise + except ValueError as e: + error_msg = f"Invalid image ID format: {image_id}" + logger.error(error_msg) + return None except Exception as e: error_msg = f"Error fetching image info: {e}" logger.error(error_msg) diff --git a/tests/services/test_civitai_client.py b/tests/services/test_civitai_client.py index ab3a5ef9..f9b2ee7e 100644 --- a/tests/services/test_civitai_client.py +++ b/tests/services/test_civitai_client.py @@ -484,9 +484,11 @@ async def test_get_model_version_info_success(monkeypatch, downloader): assert result["images"][0]["meta"]["other"] == "keep" -async def test_get_image_info_returns_first_item(monkeypatch, downloader): +async def test_get_image_info_returns_matching_item(monkeypatch, downloader): + """When API returns multiple items, return the one matching the requested ID.""" async def fake_make_request(method, url, use_auth=True, **kwargs): - return True, {"items": [{"id": 1}, {"id": 2}]} + # Requested ID is 42, but it's the second item in the response + return True, {"items": [{"id": 41}, {"id": 42, "name": "target"}, {"id": 43}]} downloader.make_request = fake_make_request @@ -494,7 +496,25 @@ async def test_get_image_info_returns_first_item(monkeypatch, downloader): result = await client.get_image_info("42") - assert result == {"id": 1} + assert result == {"id": 42, "name": "target"} + + +async def test_get_image_info_returns_none_when_id_mismatch(monkeypatch, downloader, caplog): + """When API returns items but none match the requested ID, return None and log warning.""" + async def fake_make_request(method, url, use_auth=True, **kwargs): + # Requested ID is 999, but API returns different IDs (simulating deleted/hidden image) + return True, {"items": [{"id": 1}, {"id": 2}, {"id": 3}]} + + downloader.make_request = fake_make_request + + client = await CivitaiClient.get_instance() + + result = await client.get_image_info("999") + + assert result is None + # Verify warning was logged + assert "CivitAI API returned no matching image for requested ID 999" in caplog.text + assert "Returned 3 item(s) with IDs: [1, 2, 3]" in caplog.text async def test_get_image_info_handles_missing(monkeypatch, downloader): @@ -508,3 +528,13 @@ async def test_get_image_info_handles_missing(monkeypatch, downloader): result = await client.get_image_info("42") assert result is None + + +async def test_get_image_info_handles_invalid_id(monkeypatch, downloader, caplog): + """When given a non-numeric image ID, return None and log error.""" + client = await CivitaiClient.get_instance() + + result = await client.get_image_info("not-a-number") + + assert result is None + assert "Invalid image ID format" in caplog.text