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 <noreply@anthropic.com>
This commit is contained in:
Will Miao
2026-03-25 20:33:43 +08:00
parent 6f2a01dc86
commit de3d0571f8
2 changed files with 59 additions and 6 deletions

View File

@@ -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:
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 result["items"][0]
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)

View File

@@ -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