basicmachines-co
diff --git a/‎src/basic_memory/schemas/document.py‎
Lines changed: 19 additions & 6 deletions b/‎src/basic_memory/schemas/document.py‎
Lines changed: 19 additions & 6 deletions
diff --git a/‎test-int/test_document_page_citations.py‎
Lines changed: 143 additions & 0 deletions b/‎test-int/test_document_page_citations.py‎
Lines changed: 143 additions & 0 deletions
@@ -582,10 +582,19 @@ def reject_unstructured_semantics(cls, value: str) -> str:
     @model_validator(mode="after")
     def require_exact_assembled_semantics(self) -> "DocumentAgentOutputV1":
         if any(observation.locator is not None for observation in self.observations):
-            # Trusted code owns these definitions. An agent-supplied definition
-            # must not redirect a generated citation to another resource.
-            if re.search(r"(?m)^[ \t]*\[\^document-page-[0-9]+\]:", self.body):
-                raise ValueError("agent body cannot define generated document-page citations")
+            # Reserve both references and definitions: otherwise undeclared text
+            # can borrow a generated citation without supplying a locator.
+            agent_text = [self.title, self.body]
+            for observation in self.observations:
+                agent_text.extend((observation.content, observation.context or ""))
+            for relation in self.relations:
+                agent_text.extend((relation.target, relation.context or ""))
+            if any(
+                re.search(r"\[\^document-page-[0-9]+\]", text, re.IGNORECASE) for text in agent_text
+            ):
+                raise ValueError(
+                    "agent text cannot define generated document-page citations or references"
+                )
         parsed = _parse_agent_semantics(_assemble_agent_body(self))
         expected_observations = [
             _parse_agent_semantics(_format_agent_observation(observation)).observations[0]
@@ -768,8 +777,12 @@ def enrich_document_markdown(
             raw_frontmatter.source, raw_frontmatter.extraction, observation.locator
         )
         existing = citations.get(citation.id)
-        if existing is not None and existing != citation:
-            raise ValueError("observations citing the same page must agree on its page label")
+        if existing is not None and existing.locator.page_label is not None:
+            if citation.locator.page_label not in {None, existing.locator.page_label}:
+                raise ValueError("observations citing the same page must agree on its page label")
+            # An omitted label contributes no conflicting information. Keep the
+            # explicit label regardless of observation order.
+            continue
         citations[citation.id] = citation
     citation_sources = tuple(citations.values())
     return DocumentMarkdownV1(
 
@@ -26,6 +26,7 @@
 from basic_memory.schemas.document import (
     DocumentAgentObservationV1,
     DocumentAgentOutputV1,
+    DocumentAgentRelationV1,
     DocumentIngestionStage,
     DocumentIngestionV1,
     DocumentMarkdownV1,
@@ -309,3 +310,145 @@ def test_pdf_citation_rejects_source_entry_that_disagrees_with_provenance(
     assert redirected != markdown
     with pytest.raises(ValidationError, match="trusted source PDF page"):
         parse_document_markdown(redirected)
+
+
+def test_relation_context_cannot_borrow_declared_page_citation() -> None:
+    with pytest.raises(ValidationError, match="cannot define generated"):
+        DocumentAgentOutputV1(
+            title="Citations",
+            body="",
+            observations=(
+                DocumentAgentObservationV1(
+                    category="fact", content="Claim.", locator=DocumentPageLocatorV1(page=2)
+                ),
+            ),
+            relations=(
+                DocumentAgentRelationV1(
+                    relation_type="supports",
+                    target="Evidence",
+                    context="Borrowed [^document-page-2]",
+                ),
+            ),
+        )
+
+
+def test_pdf_citation_rejects_duplicate_source_ids(raw_pdf: DocumentMarkdownV1) -> None:
+    enriched = enrich(
+        raw_pdf,
+        (
+            DocumentAgentObservationV1(
+                category="fact", content="Claim.", locator=DocumentPageLocatorV1(page=2)
+            ),
+        ),
+    )
+    assert enriched.frontmatter.sources is not None
+    duplicate = enriched.model_copy(
+        update={
+            "frontmatter": enriched.frontmatter.model_copy(
+                update={"sources": enriched.frontmatter.sources * 2}
+            )
+        }
+    )
+    with pytest.raises(ValidationError, match="citation source IDs must be unique"):
+        parse_document_markdown(assemble_document_markdown(duplicate))
+
+
+def test_page_citation_rejects_non_pdf_source(raw_pdf: DocumentMarkdownV1) -> None:
+    markdown = assemble_document_markdown(raw_pdf).replace(
+        "media_type: application/pdf", "media_type: text/plain"
+    )
+    non_pdf = parse_document_markdown(markdown)
+    with pytest.raises(ValueError, match="page citations require a PDF source"):
+        enrich(
+            non_pdf,
+            (
+                DocumentAgentObservationV1(
+                    category="fact", content="Claim.", locator=DocumentPageLocatorV1(page=2)
+                ),
+            ),
+        )
+
+
+def test_page_label_is_retained_when_later_observation_omits_it(
+    raw_pdf: DocumentMarkdownV1,
+) -> None:
+    enriched = enrich(
+        raw_pdf,
+        (
+            DocumentAgentObservationV1(
+                category="fact",
+                content="First.",
+                locator=DocumentPageLocatorV1(page=2, page_label="iv"),
+            ),
+            DocumentAgentObservationV1(
+                category="fact", content="Second.", locator=DocumentPageLocatorV1(page=2)
+            ),
+        ),
+    )
+    parsed = parse_document_markdown(assemble_document_markdown(enriched))
+    assert parsed.frontmatter.sources is not None
+    assert len(parsed.frontmatter.sources) == 1
+    assert parsed.frontmatter.sources[0].locator.page_label == "iv"
+
+
+def test_page_label_is_added_when_earlier_observation_omits_it(raw_pdf: DocumentMarkdownV1) -> None:
+    enriched = enrich(
+        raw_pdf,
+        (
+            DocumentAgentObservationV1(
+                category="fact", content="First.", locator=DocumentPageLocatorV1(page=2)
+            ),
+            DocumentAgentObservationV1(
+                category="fact",
+                content="Second.",
+                locator=DocumentPageLocatorV1(page=2, page_label="iv"),
+            ),
+        ),
+    )
+    parsed = parse_document_markdown(assemble_document_markdown(enriched))
+    assert parsed.frontmatter.sources is not None
+    assert len(parsed.frontmatter.sources) == 1
+    assert parsed.frontmatter.sources[0].locator.page_label == "iv"
+
+
+def test_agent_body_cannot_borrow_declared_page_citation() -> None:
+    with pytest.raises(ValidationError, match="cannot define generated"):
+        DocumentAgentOutputV1(
+            title="Citations",
+            body="Unsupported claim [^document-page-2]",
+            observations=(
+                DocumentAgentObservationV1(
+                    category="fact", content="Claim.", locator=DocumentPageLocatorV1(page=2)
+                ),
+            ),
+        )
+
+
+def test_unlocated_observation_cannot_borrow_declared_page_citation() -> None:
+    with pytest.raises(ValidationError, match="cannot define generated"):
+        DocumentAgentOutputV1(
+            title="Citations",
+            body="",
+            observations=(
+                DocumentAgentObservationV1(
+                    category="fact", content="Claim.", locator=DocumentPageLocatorV1(page=2)
+                ),
+                DocumentAgentObservationV1(category="fact", content="Unlocated [^document-page-2]"),
+            ),
+        )
+
+
+def test_observation_context_cannot_borrow_declared_page_citation() -> None:
+    with pytest.raises(ValidationError, match="cannot define generated"):
+        DocumentAgentOutputV1(
+            title="Citations",
+            body="",
+            observations=(
+                DocumentAgentObservationV1(
+                    category="fact",
+                    content="Claim.",
+                    locator=DocumentPageLocatorV1(page=2),
+                    context="Borrowed [^document-page-2]",
+                ),
+            ),
+        )