From b9216392c94b2fedac43e30b8fb202b3743ef011 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 13 Sep 2026 20:58:45 +0800 Subject: [PATCH 1/3] fix(retrieval): skip map-unit count equality on asset filters Classic image/table queries compared full-revision index.unit_count to the filtered unit subset, so a complete index looked unusable and returned 500. Co-authored-by: Cursor --- ...est_retrieval_classic_map_unit_contract.py | 246 ++++++++++++++++++ .../retrieval/search/map_unit_discovery.py | 29 ++- 2 files changed, 271 insertions(+), 4 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index d42148d8..57cc5077 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -781,6 +781,252 @@ async def test_classic_route_image_filter_scores_only_units_with_images( assert results[0]["chunk_type"] == "image" +async def test_classic_route_mixed_section_asset_filters_do_not_mark_index_unusable( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-mixed-{identifier}" + async with developer_api_client_factory() as api_client: + mixed = await _publish_document( + namespace=namespace, + source_file_name="mixed-assets.pdf", + chunks=[ + { + "chunk_id": f"plain-{identifier}", + "type": "text", + "content": "mixedsection plaintext filler with no assets", + "path": "mixed-assets.pdf/Root/Plain/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"chart-body-{identifier}", + "type": "text", + "content": "mixedsection chartmarker beside a plot", + "path": "mixed-assets.pdf/Root/Chart/body", + "order": 2, + "metadata": {"connect_to": [{"target": f"chart-{identifier}"}]}, + }, + { + "chunk_id": f"chart-{identifier}", + "type": "image", + "content": "mixedsection chartmarker plot", + "path": "images/mixed-chart.png", + "order": 3, + "file_path": "images/mixed-chart.png", + "metadata": {}, + }, + { + "chunk_id": f"grid-body-{identifier}", + "type": "text", + "content": "mixedsection tablemarker beside a grid", + "path": "mixed-assets.pdf/Root/Grid/body", + "order": 4, + "metadata": {"connect_to": [{"target": f"grid-{identifier}"}]}, + }, + { + "chunk_id": f"grid-{identifier}", + "type": "table", + "content": "
mixedsection tablemarker grid
", + "path": "tables/mixed-grid.html", + "order": 5, + "file_path": "tables/mixed-grid.html", + "metadata": {}, + }, + ], + ) + for extra_kind, extra_type, extra_path, extra_content in ( + ("photo", "image", "images/extra-photo.png", "unrelated landscape photo"), + ("diagram", "image", "images/extra-diagram.png", "unrelated diagram"), + ("sheet", "table", "tables/extra-sheet.html", "
unrelated sheet
"), + ("grid", "table", "tables/extra-grid.html", "
unrelated grid
"), + ): + extra_chunk_id = f"{extra_kind}-{identifier}" + await _publish_document( + namespace=namespace, + source_file_name=f"extra-{extra_kind}.pdf", + chunks=[ + { + "chunk_id": f"{extra_kind}-body-{identifier}", + "type": "text", + "content": f"unrelated {extra_kind} caption", + "path": f"extra-{extra_kind}.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": extra_chunk_id}]}, + }, + { + "chunk_id": extra_chunk_id, + "type": extra_type, + "content": extra_content, + "path": extra_path, + "order": 2, + "file_path": extra_path, + "metadata": {}, + }, + ], + ) + async with contract_db_session() as db: + mixed_units = list( + ( + await db.execute( + select(DocumentMapUnit).where( + DocumentMapUnit.document_id == mixed["document_id"] + ) + ) + ).scalars() + ) + assert len(mixed_units) >= 3 + assert any(unit.has_image and not unit.has_table for unit in mixed_units) + assert any(unit.has_table and not unit.has_image for unit in mixed_units) + assert any(not unit.has_image and not unit.has_table for unit in mixed_units) + + image_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "mixedsection chartmarker", + "top_k": 1, + "use_agentic": False, + "chunk_types": ["image"], + }, + ) + table_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "mixedsection tablemarker", + "top_k": 1, + "use_agentic": False, + "chunk_types": ["table"], + }, + ) + unfiltered_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "mixedsection plaintext", + "top_k": 1, + "use_agentic": False, + }, + ) + + assert image_response.status_code == 200 + image_body = cast(dict[str, object], image_response.json()) + image_results = cast(list[dict[str, object]], image_body["results"]) + assert image_body["router_used"] == "classic_topk" + assert [row["chunk_id"] for row in image_results] == [f"chart-{identifier}"] + assert image_results[0]["chunk_type"] == "image" + + assert table_response.status_code == 200 + table_body = cast(dict[str, object], table_response.json()) + table_results = cast(list[dict[str, object]], table_body["results"]) + assert table_body["router_used"] == "classic_topk" + assert [row["chunk_id"] for row in table_results] == [f"grid-{identifier}"] + assert table_results[0]["chunk_type"] == "table" + + assert unfiltered_response.status_code == 200 + unfiltered_body = cast(dict[str, object], unfiltered_response.json()) + assert unfiltered_body["router_used"] == "classic_topk" + + +@pytest.mark.parametrize( + "incomplete_index_kind", ["legacy_format", "missing_index"] +) +async def test_classic_route_image_filter_raises_when_index_is_unusable( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + incomplete_index_kind: str, +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-image-unusable-{incomplete_index_kind}-{identifier}" + async with developer_api_client_factory() as api_client: + document = await _publish_document( + namespace=namespace, + source_file_name="unusable-image.pdf", + chunks=[ + { + "chunk_id": f"plain-{identifier}", + "type": "text", + "content": "unusable image plaintext filler", + "path": "unusable-image.pdf/Root/Plain/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"body-{identifier}", + "type": "text", + "content": "unusable image marker next to a chart", + "path": "unusable-image.pdf/Root/Chart/body", + "order": 2, + "metadata": {"connect_to": [{"target": f"chart-{identifier}"}]}, + }, + { + "chunk_id": f"chart-{identifier}", + "type": "image", + "content": "unusable image marker chart", + "path": "images/unusable-chart.png", + "order": 3, + "file_path": "images/unusable-chart.png", + "metadata": {}, + }, + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="extra-image.pdf", + chunks=[ + { + "chunk_id": f"extra-body-{identifier}", + "type": "text", + "content": "unrelated extra caption", + "path": "extra-image.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": f"extra-{identifier}"}]}, + }, + { + "chunk_id": f"extra-{identifier}", + "type": "image", + "content": "unrelated extra photo", + "path": "images/extra.png", + "order": 2, + "file_path": "images/extra.png", + "metadata": {}, + }, + ], + ) + if incomplete_index_kind == "legacy_format": + await ContractDatabase.execute( + """ + UPDATE document_map_unit_indexes + SET format_version = 1 + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, + ) + else: + await ContractDatabase.execute( + """ + DELETE FROM document_map_unit_indexes + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, + ) + with pytest.raises(RuntimeError, match="map-unit index is incomplete"): + await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unusable image marker", + "top_k": 1, + "use_agentic": False, + "chunk_types": ["image"], + }, + ) + + async def test_connected_hydration_does_not_load_legacy_job_chunks( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 9a1561c9..e9510286 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -451,9 +451,16 @@ async def map_unit_discovery( has_index_storage_mismatch = indexed_unit_count != int( actual_unit_count or 0 ) or indexed_token_count != int(actual_token_count or 0) + # Equality against indexes.unit_count is only valid when unit_rows is the + # full inventory of those revisions. Image/table type filters, signal + # paths, and exclude_sections load a subset, so require the looser + # undercount check used for unfiltered token projection. + unit_rows_are_full_revision_inventory = not ( + type_clause or signal_paths or exclude_sections + ) has_index_unit_count_mismatch = ( indexed_unit_count < len(unit_rows) - if is_unfiltered_scope + if is_unfiltered_scope or not unit_rows_are_full_revision_inventory else indexed_unit_count != len(unit_rows) ) is_index_format_incompatible = any( @@ -488,8 +495,9 @@ async def map_unit_discovery( _token_count, ) in index_parts ) + has_revision_coverage_mismatch = len(index_parts) != len(expected_revisions) has_unusable_index = ( - len(index_parts) != len(expected_revisions) + has_revision_coverage_mismatch or has_index_unit_count_mismatch or has_index_storage_mismatch or is_index_format_incompatible @@ -505,11 +513,24 @@ async def map_unit_discovery( ) except Exception as exc: logger.warning("retrieval index readiness publish failed: %s", exc) + if has_revision_coverage_mismatch: + unusable_reason = "revision_coverage" + elif is_index_format_incompatible: + unusable_reason = "format" + elif has_index_storage_mismatch: + unusable_reason = "storage" + else: + unusable_reason = "unit_count_mismatch" + chunk_types_label = ( + ",".join(sorted(chunk_types)) if chunk_types else "none" + ) raise RuntimeError( "retrieval map-unit index is incomplete or incompatible " - f"(user_id={user_id} namespace={namespace} " + f"(reason={unusable_reason} user_id={user_id} namespace={namespace} " f"expected_revisions={len(expected_revisions)} " - f"indexed_revisions={len(index_parts)})" + f"indexed_revisions={len(index_parts)} " + f"indexed_unit_count={indexed_unit_count} unit_rows={len(unit_rows)} " + f"unfiltered={is_unfiltered_scope} chunk_types={chunk_types_label})" ) if has_incomplete_index_statistics: logger.warning( From 6ff0a09c7fdb51f528bad4b2f0c5449ff5cecf3d Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 13 Sep 2026 21:06:25 +0800 Subject: [PATCH 2/3] Potential fix for pull request finding 'CodeQL / Potentially uninitialized local variable' Initialize the revision-inventory flag before the async storage-mismatch branch so CodeQL sees it on every path. Co-authored-by: Cursor --- .../retrieval/search/map_unit_discovery.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index e9510286..b65b42c6 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -417,6 +417,13 @@ async def map_unit_discovery( indexed_unit_count = sum( unit_count for _path_idf, _content_idf, unit_count, *_rest in index_parts ) + # Equality against indexes.unit_count is only valid when unit_rows is the + # full inventory of those revisions. Image/table type filters, signal + # paths, and exclude_sections load a subset, so require the looser + # undercount check used for unfiltered token projection. + unit_rows_are_full_revision_inventory: bool = not ( + type_clause or signal_paths or exclude_sections + ) has_index_storage_mismatch: bool = False if is_unfiltered_scope and not unit_rows and expected_revisions: storage_counts: tuple[int | None, int | None] = cast( @@ -451,13 +458,6 @@ async def map_unit_discovery( has_index_storage_mismatch = indexed_unit_count != int( actual_unit_count or 0 ) or indexed_token_count != int(actual_token_count or 0) - # Equality against indexes.unit_count is only valid when unit_rows is the - # full inventory of those revisions. Image/table type filters, signal - # paths, and exclude_sections load a subset, so require the looser - # undercount check used for unfiltered token projection. - unit_rows_are_full_revision_inventory = not ( - type_clause or signal_paths or exclude_sections - ) has_index_unit_count_mismatch = ( indexed_unit_count < len(unit_rows) if is_unfiltered_scope or not unit_rows_are_full_revision_inventory From 2118ac120869ed67b43f754ca396d9d17cc326a9 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 13 Sep 2026 21:12:23 +0800 Subject: [PATCH 3/3] Potential fix for pull request finding 'CodeQL / Potentially uninitialized local variable' Inline the subset-inventory condition so the extra local is gone. CodeQL kept flagging the use of that name even after it was assigned earlier. Co-authored-by: Cursor --- .../services/retrieval/search/map_unit_discovery.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index b65b42c6..790dc92d 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -417,13 +417,6 @@ async def map_unit_discovery( indexed_unit_count = sum( unit_count for _path_idf, _content_idf, unit_count, *_rest in index_parts ) - # Equality against indexes.unit_count is only valid when unit_rows is the - # full inventory of those revisions. Image/table type filters, signal - # paths, and exclude_sections load a subset, so require the looser - # undercount check used for unfiltered token projection. - unit_rows_are_full_revision_inventory: bool = not ( - type_clause or signal_paths or exclude_sections - ) has_index_storage_mismatch: bool = False if is_unfiltered_scope and not unit_rows and expected_revisions: storage_counts: tuple[int | None, int | None] = cast( @@ -458,9 +451,13 @@ async def map_unit_discovery( has_index_storage_mismatch = indexed_unit_count != int( actual_unit_count or 0 ) or indexed_token_count != int(actual_token_count or 0) + # Equality against indexes.unit_count is only valid when unit_rows is the + # full inventory of those revisions. Image/table type filters, signal + # paths, and exclude_sections load a subset, so require the looser + # undercount check used for unfiltered token projection. has_index_unit_count_mismatch = ( indexed_unit_count < len(unit_rows) - if is_unfiltered_scope or not unit_rows_are_full_revision_inventory + if is_unfiltered_scope or type_clause or signal_paths or exclude_sections else indexed_unit_count != len(unit_rows) ) is_index_format_incompatible = any(