Coverage for src/lilbee/server/handlers/documents.py: 100%
65 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Document listing, deletion, and source-content handlers."""
3from __future__ import annotations
5import asyncio
6import mimetypes
8from litestar.exceptions import ClientException
9from litestar.status_codes import HTTP_409_CONFLICT
11from lilbee.app.services import get_services
12from lilbee.app.titles import stored_web_titles
13from lilbee.data.store import DOCUMENTS_PAGE_DEFAULT
14from lilbee.server.models import (
15 DocumentInfo,
16 DocumentListResponse,
17 DocumentRemoveResponse,
18 SourceContentResponse,
19)
21# Windows mimetypes reads from the registry, which may not define ``.md``
22# as ``text/markdown``. Pin the mapping at import time; ``add_type`` is
23# idempotent so repeated imports are safe.
24mimetypes.add_type("text/markdown", ".md")
27# Types that can carry script even within an "inline-rendered" category.
28# Keep the deny narrow and explicit. Broadening this set is a security-relevant
29# change: file an issue with the ``security`` label before adding entries.
30_RAW_INLINE_RENDER_DENY: frozenset[str] = frozenset(
31 {
32 "text/html",
33 "text/javascript",
34 "application/javascript",
35 "application/xhtml+xml",
36 "text/css",
37 "image/svg+xml",
38 # An xml-stylesheet PI can pull in XSLT that emits script. Both
39 # spellings, because which one a ``.xml`` file resolves to depends on
40 # the host mimetypes database.
41 "text/xml",
42 "application/xml",
43 }
44)
47def _is_safe_for_inline_render(content_type: str) -> bool:
48 """Whether ``raw=1`` may serve this Content-Type as-is.
50 Trusted categories (``text/*``, ``image/*``, ``application/pdf``) pass
51 through, with named exceptions for types that embed executable script.
52 Everything else degrades to ``application/octet-stream`` so an attacker-
53 renamed file (e.g. ``evil.html``) cannot trick a browser into rendering
54 it inline within the plugin origin.
55 """
56 if content_type in _RAW_INLINE_RENDER_DENY:
57 return False
58 if content_type == "application/pdf":
59 return True
60 return content_type.startswith("text/") or content_type.startswith("image/")
63def _imported_source_markdown(source: str) -> str | None:
64 """Page texts joined in page order; ``None`` when the source has none."""
65 rows = get_services().store.get_page_texts(source)
66 if not rows:
67 return None
68 ordered = sorted(rows, key=lambda row: row["page"])
69 return "\n\n".join(row["text"] for row in ordered)
72async def delete_documents(names: list[str]) -> DocumentRemoveResponse:
73 """Remove documents by source name, folder, or glob (source files are kept)."""
74 from lilbee.app.ingest import LIBRARY_REFUSALS, remove_documents_durably
76 # Deletes store rows and takes the wiki build mutex to drop the removed
77 # documents from the browse index, so it cannot run on the event loop: a
78 # build in flight holds that mutex for the length of the whole run.
79 try:
80 result = await asyncio.to_thread(remove_documents_durably, names)
81 except LIBRARY_REFUSALS as exc:
82 raise ClientException(detail=str(exc), status_code=HTTP_409_CONFLICT) from exc
83 return DocumentRemoveResponse(removed=result.removed, not_found=result.not_found)
86async def list_documents(
87 search: str = "",
88 limit: int = DOCUMENTS_PAGE_DEFAULT,
89 offset: int = 0,
90) -> DocumentListResponse:
91 """Return indexed documents with metadata, paginated and filterable.
93 Pagination and the filename filter are pushed into LanceDB via
94 ``Store.get_sources(search=..., limit=..., offset=...)`` and the
95 total comes from ``Store.count_sources(search=...)`` so neither
96 call materializes the full SOURCES table per request.
97 """
98 store = get_services().store
99 search_term = search or None
100 page = store.get_sources(search=search_term, limit=limit, offset=offset)
101 total = store.count_sources(search=search_term)
102 return DocumentListResponse(
103 documents=[
104 DocumentInfo(
105 filename=s["filename"],
106 chunk_count=s.get("chunk_count", 0),
107 ingested_at=s.get("ingested_at", ""),
108 )
109 for s in page
110 ],
111 total=total,
112 limit=limit,
113 offset=offset,
114 # Gated on a non-empty page: a concurrent writer shrinking SOURCES
115 # between count and fetch leaves a stale total, and a client reading
116 # has_more would spin past the end.
117 has_more=len(page) > 0 and (offset + len(page)) < total,
118 )
121async def get_source_content(
122 source: str, raw: bool = False
123) -> SourceContentResponse | tuple[bytes, str]:
124 """Return a stored source file: JSON with markdown text for text types, or
125 ``(bytes, content_type)`` when *raw* is True. Binary types return empty
126 markdown so clients know to re-request with ``raw=1``.
128 Reads the file off the event loop so a large source doesn't stall it.
129 """
130 return await asyncio.to_thread(_get_source_content_sync, source, raw)
133def _get_source_content_sync(source: str, raw: bool) -> SourceContentResponse | tuple[bytes, str]:
134 """Blocking body of :func:`get_source_content`: path validation + file read."""
135 if not source or not source.strip():
136 raise ValueError("source must not be empty")
137 from lilbee.data.ingest.discovery import resolve_source_path_checked
139 resolved = resolve_source_path_checked(source)
140 if resolved is None:
141 raise ValueError(f"source path escapes its root: {source}")
142 if not resolved.is_file():
143 # Imported sources have no file on disk; their text lives in the page-text store.
144 markdown = _imported_source_markdown(source)
145 if markdown is None:
146 raise FileNotFoundError(source)
147 if raw:
148 return markdown.encode("utf-8"), "text/markdown"
149 return SourceContentResponse(
150 markdown=markdown, content_type="text/markdown", title=_source_title(source, markdown)
151 )
153 content_type, _ = mimetypes.guess_type(resolved.name)
154 if content_type is None:
155 content_type = "application/octet-stream"
157 if raw:
158 # Cap raw responses to inline-render-safe categories; anything else
159 # degrades to a binary download so attacker-renamed files (e.g.
160 # evil.html) can't trick the embedding browser into running script
161 # under our origin.
162 served_type = (
163 content_type if _is_safe_for_inline_render(content_type) else "application/octet-stream"
164 )
165 return resolved.read_bytes(), served_type
167 if not content_type.startswith("text/"):
168 return SourceContentResponse(markdown="", content_type=content_type, title=None)
170 text = resolved.read_text(encoding="utf-8", errors="replace")
171 return SourceContentResponse(
172 markdown=text, content_type=content_type, title=_source_title(source, text)
173 )
176def _source_title(source: str, text: str) -> str | None:
177 """A crawled page's stored title, else the title the text gives itself, else None."""
178 from lilbee.wiki.index import parse_title
180 return stored_web_titles([source]).get(source) or parse_title(text) or None