diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 2038a8d..4f18071 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -326,6 +326,64 @@ def _extract_civitai_red_candidates(html: str, limit: int = 5) -> List[Dict[str, return candidates +_DEVALUE_UNDEFINED = object() + + +def _decode_devalue_trpc_data(value: str) -> Optional[Dict[str, Any]]: + """Decode CivitAI's alternate devalue-serialized tRPC data payload.""" + try: + parsed = json.loads(value) + except (TypeError, ValueError): + return None + + if isinstance(parsed, dict): + return parsed + if not isinstance(parsed, list) or not parsed: + return None + + memo: Dict[int, Any] = {} + + def hydrate_reference(reference: Any) -> Any: + if isinstance(reference, bool) or reference is None: + return reference + if not isinstance(reference, int): + return reference + if reference < 0: + return _DEVALUE_UNDEFINED + if reference >= len(parsed): + return _DEVALUE_UNDEFINED + return hydrate_index(reference) + + def hydrate_index(index: int) -> Any: + if index in memo: + return memo[index] + + raw_value = parsed[index] + if isinstance(raw_value, dict): + hydrated: Any = {} + memo[index] = hydrated + for key, reference in raw_value.items(): + resolved = hydrate_reference(reference) + if resolved is not _DEVALUE_UNDEFINED: + hydrated[key] = resolved + return hydrated + + if isinstance(raw_value, list): + hydrated = [] + memo[index] = hydrated + for reference in raw_value: + resolved = hydrate_reference(reference) + if resolved is not _DEVALUE_UNDEFINED: + hydrated.append(resolved) + return hydrated + + memo[index] = raw_value + return raw_value + + decoded = hydrate_index(0) + return decoded if isinstance(decoded, dict) else None + + def _extract_trpc_model_candidates( payload: Any, limit: int = 5 ) -> List[Dict[str, Optional[int]]]: @@ -333,12 +391,26 @@ def _extract_trpc_model_candidates( candidates: List[Dict[str, Optional[int]]] = [] seen = set() - items = ( - payload.get("result", {}) - .get("data", {}) - .get("json", {}) - .get("items", []) - ) + if not isinstance(payload, dict): + return candidates + + result = payload.get("result") + if not isinstance(result, dict): + return candidates + + data = result.get("data") + if isinstance(data, str): + data = _decode_devalue_trpc_data(data) + if not isinstance(data, dict): + return candidates + + json_data = data.get("json", data) + if isinstance(json_data, str): + json_data = _decode_devalue_trpc_data(json_data) + if not isinstance(json_data, dict): + return candidates + + items = json_data.get("items", []) if not isinstance(items, list): return candidates diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 8b7de44..a514a7b 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -46,6 +46,7 @@ from core.sources.civitai import ( _enrich_model_info_with_details, _find_matching_file_in_versions, _extract_public_api_model_candidates, + _extract_trpc_model_candidates, _search_civitai_public_api_candidates, _search_civitai_red_candidates, _search_civitai_trpc_candidates, @@ -634,6 +635,66 @@ class CivitaiRedSearchTests(unittest.TestCase): class CivitaiTrpcSearchTests(unittest.TestCase): + def test_extracts_candidates_from_devalue_serialized_data(self): + payload = { + "result": { + "data": json.dumps( + [ + {"items": 1, "nextCursor": -1}, + [2], + {"id": 3, "version": 4}, + 2725430, + {"id": 5}, + 3071760, + ], + separators=(",", ":"), + ) + } + } + + result = _extract_trpc_model_candidates(payload) + + self.assertEqual( + result, + [{"model_id": 2725430, "version_id": 3071760}], + ) + + def test_unrecognized_trpc_shapes_return_no_candidates(self): + payloads = [ + None, + "not-an-object", + {"result": "not-an-object"}, + {"result": {"data": "not-json"}}, + {"result": {"data": {"json": "not-json"}}}, + {"result": {"data": {"json": {"items": "not-a-list"}}}}, + ] + + for payload in payloads: + with self.subTest(payload=payload): + self.assertEqual(_extract_trpc_model_candidates(payload), []) + + @patch("core.sources.civitai.requests.get") + def test_devalue_serialized_empty_items_does_not_abort_search(self, mock_get): + response = MagicMock() + response.status_code = 200 + response.headers = {"content-type": "application/json"} + response.text = ( + '{"result":{"data":"[{\\"items\\":1,\\"nextCursor\\":-1},[]]"}}' + ) + response.json.return_value = { + "result": {"data": '[{"items":1,"nextCursor":-1},[]]'} + } + mock_get.return_value = response + + result = _search_civitai_trpc_candidates( + "KNPV3_1.safetensors", + model_type="loras", + limit=5, + ) + + self.assertEqual(result, []) + mock_get.assert_called_once() + @patch("core.sources.civitai.requests.get") def test_does_not_retry_without_type_filter_when_typed_search_is_empty(self, mock_get): empty_response = MagicMock()