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