RAGcore-Public/tests/security/citations/test_citations.py

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,
            }
        )