from __future__ import annotations from datetime import UTC, datetime from uuid import uuid4 import pytest from ragcore.application.citations import CitationService from ragcore.application.retrieval.authorization import AuthorizedRetrievalScope from ragcore.domain.authorization import Principal, PrincipalType, ResourceNotFound from ragcore.domain.citations import Citation from ragcore.domain.control import Role def citation(*, workspace_id, space_id, version_id=None) -> Citation: return Citation.create( workspace_id=workspace_id, space_id=space_id, document_id=uuid4(), document_version_id=version_id or uuid4(), chunk_id=uuid4(), parent_chunk_id=None, location={"page": 3, "line_start": 10, "line_end": 14}, excerpt="A stable cited excerpt.", source_uri="https://example.test/runbook.md", acquired_at=datetime(2026, 8, 1, 12, 0, tzinfo=UTC), retrieval_run_id=uuid4(), index_alias="rag_dense_hybrid_active", profile_version="hybrid-v1", ) class InMemoryCitationRepository: def __init__(self) -> None: self.rows: dict[object, Citation] = {} async def create(self, value: Citation) -> Citation: self.rows.setdefault(value.id, value) return self.rows[value.id] async def get_authorized( self, citation_id, *, workspace_id, effective_space_ids, ) -> Citation | None: value = self.rows.get(citation_id) if ( value is None or value.workspace_id != workspace_id or value.space_id not in effective_space_ids ): return None return value def principal(workspace_id) -> Principal: return Principal( id=uuid4(), workspace_id=workspace_id, principal_type=PrincipalType.HUMAN, roles=frozenset({Role.VIEWER}), ) @pytest.mark.asyncio async def test_citation_survives_document_supersession_and_creation_is_idempotent() -> None: workspace_id = uuid4() space_id = uuid4() scope = AuthorizedRetrievalScope(workspace_id, frozenset({space_id})) repository = InMemoryCitationRepository() service = CitationService(repository) first = citation(workspace_id=workspace_id, space_id=space_id) stored_first = await service.create(principal(workspace_id), scope, first) stored_again = await service.create(principal(workspace_id), scope, first) superseding = citation( workspace_id=workspace_id, space_id=space_id, version_id=uuid4(), ) await service.create(principal(workspace_id), scope, superseding) resolved = await service.resolve(principal(workspace_id), scope, first.id) assert stored_first == stored_again == first assert resolved == first assert len(repository.rows) == 2 @pytest.mark.asyncio async def test_guessed_or_forbidden_citation_ids_are_hidden() -> None: workspace_id = uuid4() allowed_space = uuid4() forbidden_space = uuid4() scope = AuthorizedRetrievalScope(workspace_id, frozenset({allowed_space})) repository = InMemoryCitationRepository() service = CitationService(repository) allowed_principal = principal(workspace_id) forbidden = citation(workspace_id=workspace_id, space_id=forbidden_space) repository.rows[forbidden.id] = forbidden with pytest.raises(ResourceNotFound): await service.resolve(allowed_principal, scope, forbidden.id) with pytest.raises(ResourceNotFound): await service.resolve(allowed_principal, scope, uuid4()) with pytest.raises(ResourceNotFound): await service.create(allowed_principal, scope, forbidden) def test_excerpt_checksum_and_location_are_stable_and_tampering_fails() -> None: workspace_id = uuid4() space_id = uuid4() inputs = { "workspace_id": workspace_id, "space_id": space_id, "document_id": uuid4(), "document_version_id": uuid4(), "chunk_id": uuid4(), "parent_chunk_id": None, "location": {"line_end": 14, "line_start": 10, "page": 3}, "excerpt": "A stable cited excerpt.", "source_uri": "https://example.test/runbook.md", "acquired_at": datetime(2026, 8, 1, 12, 0, tzinfo=UTC), "retrieval_run_id": uuid4(), "index_alias": "rag_dense_hybrid_active", "profile_version": "hybrid-v1", } first = Citation.create(**inputs) second = Citation.create(**inputs) assert first.id == second.id assert first.excerpt_checksum == second.excerpt_checksum assert list(first.location) == ["line_end", "line_start", "page"] with pytest.raises(ValueError, match="checksum"): Citation( **{ **inputs, "id": first.id, "excerpt_checksum": "0" * 64, } ) def test_location_and_identity_cannot_be_forged_after_creation() -> None: workspace_id = uuid4() space_id = uuid4() value = citation(workspace_id=workspace_id, space_id=space_id) with pytest.raises(TypeError): value.location["page"] = 99 # type: ignore[index] with pytest.raises(ValueError, match="id does not match"): Citation( **{ "id": uuid4(), "workspace_id": value.workspace_id, "space_id": value.space_id, "document_id": value.document_id, "document_version_id": value.document_version_id, "chunk_id": value.chunk_id, "parent_chunk_id": value.parent_chunk_id, "location": dict(value.location), "excerpt": value.excerpt, "excerpt_checksum": value.excerpt_checksum, "source_uri": value.source_uri, "acquired_at": value.acquired_at, "retrieval_run_id": value.retrieval_run_id, "index_alias": value.index_alias, "profile_version": value.profile_version, } )