diff --git a/pageindex/backend/cloud.py b/pageindex/backend/cloud.py index 6d2b6f2..bbfd863 100644 --- a/pageindex/backend/cloud.py +++ b/pageindex/backend/cloud.py @@ -239,6 +239,23 @@ class CloudBackend: raise DocumentNotFoundError(f"Document {doc_id} not found") from e raise + def _get_metadata(self, doc_id: str) -> dict: + return self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/metadata/") + + def _require_document(self, collection: str, doc_id: str) -> dict | None: + """Membership guard (the server's doc endpoints are user-scoped, not + folder-scoped, so this is checked client-side). Returns the doc's + metadata, or None when folders are unavailable on this plan.""" + folder_id = self._get_folder_id(collection) + if folder_id is None: + return None + meta = self._get_metadata(doc_id) + if meta.get("folderId") != folder_id: + raise DocumentNotFoundError( + f"Document {doc_id} not found in collection '{collection}'" + ) + return meta + def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> dict: if include_text: import warnings @@ -249,7 +266,9 @@ class CloudBackend: UserWarning, stacklevel=3, ) - resp = self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/metadata/") + resp = self._require_document(collection, doc_id) + if resp is None: + resp = self._get_metadata(doc_id) # Fetch structure in the same call via tree endpoint tree_resp = self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/", params={"type": "tree", "summary": "true"}) @@ -264,12 +283,14 @@ class CloudBackend: } def get_document_structure(self, collection: str, doc_id: str) -> list: + self._require_document(collection, doc_id) resp = self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/", params={"type": "tree", "summary": "true"}) raw_tree = resp.get("tree", resp.get("structure", resp.get("result", []))) return self._normalize_tree(raw_tree) def get_page_content(self, collection: str, doc_id: str, pages: str) -> list: + self._require_document(collection, doc_id) resp = self._doc_request(doc_id, "GET", f"/doc/{self._enc(doc_id)}/", params={"type": "ocr", "format": "page"}) # Filter to requested pages @@ -344,6 +365,7 @@ class CloudBackend: offset += page_size def delete_document(self, collection: str, doc_id: str) -> None: + self._require_document(collection, doc_id) self._doc_request(doc_id, "DELETE", f"/doc/{self._enc(doc_id)}/") # ── Query (uses cloud chat/completions, no LLM key needed) ──────────── diff --git a/pageindex/backend/protocol.py b/pageindex/backend/protocol.py index c5b390a..8ab489c 100644 --- a/pageindex/backend/protocol.py +++ b/pageindex/backend/protocol.py @@ -21,7 +21,8 @@ class Backend(Protocol): def list_collections(self) -> list[str]: ... def delete_collection(self, name: str) -> None: ... - # Document management + # Document management. Contract: a doc_id not belonging to `collection` + # must behave exactly like a missing one (DocumentNotFoundError). def add_document(self, collection: str, file_path: str) -> str: ... def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> DocumentDetail: ... def get_document_structure(self, collection: str, doc_id: str) -> list: ... diff --git a/pageindex/collection.py b/pageindex/collection.py index fe0110d..8d974d0 100644 --- a/pageindex/collection.py +++ b/pageindex/collection.py @@ -71,7 +71,8 @@ class Collection: ``include_text=True`` fills each node's text from cached pages (local backend only; can be large — avoid for LLM contexts). Raises - ``DocumentNotFoundError`` if the doc_id is unknown. + ``DocumentNotFoundError`` if the doc_id is unknown or belongs to + another collection. """ return self._backend.get_document(self._name, doc_id, include_text=include_text) @@ -91,7 +92,8 @@ class Collection: def delete_document(self, doc_id: str) -> None: """Delete a document and its stored files/artifacts. - Raises ``DocumentNotFoundError`` if the doc_id is unknown. + Raises ``DocumentNotFoundError`` if the doc_id is unknown or belongs + to another collection. """ self._backend.delete_document(self._name, doc_id) diff --git a/tests/test_cloud_backend.py b/tests/test_cloud_backend.py index 1fb0cfa..fcdb081 100644 --- a/tests/test_cloud_backend.py +++ b/tests/test_cloud_backend.py @@ -243,6 +243,7 @@ def test_doc_404_maps_to_document_not_found(monkeypatch): def test_get_document_include_text_warns(monkeypatch): backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["col"] = None # folders unavailable on this plan monkeypatch.setattr(backend, "_doc_request", lambda *a, **k: {"tree": []}) with pytest.warns(UserWarning, match="include_text is not supported"): backend.get_document("col", "d1", include_text=True) @@ -254,6 +255,7 @@ def test_get_page_content_preserves_images(monkeypatch): # shape) so the SDK-prompted UI can render figures — omitting it only when # empty. Real API uses page_index/markdown/images keys. backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["col"] = None # folders unavailable on this plan imgs = [{"path": "fig1.png", "width": 640, "height": 480}] monkeypatch.setattr(backend, "_doc_request", lambda *a, **k: {"result": [ {"page_index": 1, "markdown": "page one", "images": imgs}, @@ -266,6 +268,69 @@ def test_get_page_content_preserves_images(monkeypatch): ] +# ── collection membership guard on doc-scoped operations ──────────────────── + +def _membership_backend(monkeypatch, doc_folder="f-col"): + """Collection 'col' → folder 'f-col'; fake server holds doc 'd1' in + ``doc_folder``. Returns (backend, log of (method, path)).""" + backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["col"] = "f-col" + log = [] + + def fake_request(method, path, **kwargs): + log.append((method, path)) + if path.endswith("/metadata/"): + return {"id": "d1", "name": "doc.pdf", "folderId": doc_folder} + if method == "DELETE": + return {} + return {"tree": [], "result": []} + + monkeypatch.setattr(backend, "_request", fake_request) + return backend, log + + +def test_doc_ops_reject_doc_from_another_collection(monkeypatch): + backend, log = _membership_backend(monkeypatch, doc_folder="f-other") + with pytest.raises(DocumentNotFoundError, match="collection 'col'"): + backend.get_document("col", "d1") + with pytest.raises(DocumentNotFoundError, match="collection 'col'"): + backend.get_document_structure("col", "d1") + with pytest.raises(DocumentNotFoundError, match="collection 'col'"): + backend.get_page_content("col", "d1", "1") + with pytest.raises(DocumentNotFoundError, match="collection 'col'"): + backend.delete_document("col", "d1") + # guard must fail before any destructive/content request goes out + assert all(method == "GET" and path.endswith("/metadata/") for method, path in log) + + +def test_delete_document_scoped_to_collection(monkeypatch): + backend, log = _membership_backend(monkeypatch) + backend.delete_document("col", "d1") + assert ("DELETE", "/doc/d1/") in log + + +def test_get_document_reuses_membership_metadata(monkeypatch): + backend, log = _membership_backend(monkeypatch) + doc = backend.get_document("col", "d1") + assert doc["doc_id"] == "d1" + metadata_calls = [p for _, p in log if p.endswith("/metadata/")] + assert len(metadata_calls) == 1 + + +def test_doc_ops_skip_guard_when_folders_unavailable(monkeypatch): + backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["col"] = None + log = [] + + def fake_request(method, path, **kwargs): + log.append((method, path)) + return {} + + monkeypatch.setattr(backend, "_request", fake_request) + backend.delete_document("col", "d1") + assert log == [("DELETE", "/doc/d1/")] + + # ── query_stream: terminal contract and error propagation ─────────────────── def _collect_events(backend, **kwargs):