Studio: WIP — RAG preview UI, locator/auth refactor, tests, fixtures (pre-merge snapshot)

Snapshot taken before fast-forwarding feature/rag to origin and merging main.
Bundles in-flight work so the merge has a clean tree:

Frontend
- PDF preview panel (preview-panel, preview-pdf-view, preview-text-view,
  preview-unavailable) with lazy-rendered page thumbnail rail
- Resizable preview slot via useResizablePanelWidth hook (drag handle,
  localStorage persistence, viewport clamping)
- Neutral scrollbar + Source Excerpt card restyle (no brand-coloured rail)
- Preview-store + chat-adapter / rag-api / kb-detail wiring
- Frontend test harness (vitest.config, setupTests, biome update) and the
  paired __tests__ suites for preview, sources, document-row, chat-adapter,
  rag-api, knowledge-bases-tab, search-knowledge-base-tool-ui

Backend
- RAG locator + authorization modules with chunking / retrieval / tool /
  vector_store / studio_db updates
- Paired test_rag_* suites (authorization, locators, locator_backfill,
  locator_migration, preview_routes, preview_target_locators, source_identity)

Other
- tests/fixtures/rag-preview for preview route fixtures (sample.pdf,
  sample.txt, make_fixture_pdf.py)
- .gitignore + package(-lock).json adjustments for the new test runner

Will be squashed/reworked via interactive rebase after main is merged.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Etherll 2026-05-28 00:05:13 +03:00
commit 27b0a50a84
51 changed files with 9984 additions and 1153 deletions

3
.gitignore vendored
View file

@ -235,3 +235,6 @@ package-lock.json
!studio/backend/core/data_recipe/oxc-validator/package-lock.json
!studio/package-lock.json
llama.cpp/
/.Codex
/.gemini
/.antigravitycli

View file

@ -0,0 +1,121 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Subject-scoped authorization for RAG document preview routes.
Used by `/api/rag/documents/{document_id}/file` and
`/api/rag/documents/{document_id}/preview-target` to enforce that the
current authenticated subject is allowed to see a given document and
chunk. Existence and authorization failures collapse to a single 404
so the API does not leak document IDs to a non-owner.
"""
from __future__ import annotations
import sqlite3
from fastapi import HTTPException
from storage.studio_db import get_connection
_NOT_FOUND_DETAIL = "Document not found"
def document_for_subject_or_404(
document_id: str,
current_subject: str,
) -> sqlite3.Row:
"""Return the `rag_documents` row if `current_subject` may access it.
Authorization rules:
- KB documents: the document's KB must have
`rag_knowledge_bases.owner_user_id == current_subject`. A KB with a
NULL owner is not accessible through this helper (legacy pre-auth
rows must be migrated or accessed via admin tooling).
- Thread documents: thread-scoped RAG documents are gated by an
explicit single-user invariant for Studio's current release. The
`chat_threads` table does not yet carry an `owner_user_id` column,
so we cannot bind a thread to a specific subject in the schema.
The helper still requires (a) an authenticated subject (enforced
by the route's `Depends(get_current_subject)`) and (b) that the
referenced thread actually exists in `chat_threads`. A missing
thread row collapses to 404 so a non-existent thread cannot
silently grant access through a dangling `thread_id`.
# TODO(thread-owner): once `chat_threads.owner_user_id` exists,
# join through it the same way KB documents do and drop the
# single-user invariant. Update the test
# `tests/test_rag_authorization.py::test_thread_doc_other_user_404`
# to assert per-user isolation rather than thread existence.
Both not-found and not-authorized raise `HTTPException(404)` with the
same detail string. Callers must NOT distinguish the two cases in
their response, to avoid leaking document existence to a non-owner.
Returns the document row so the caller can read `stored_path`,
`filename`, `content_type`, etc. without re-querying.
"""
if not document_id or not current_subject:
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
with get_connection() as conn:
row = conn.execute(
"SELECT * FROM rag_documents WHERE id = ?",
(document_id,),
).fetchone()
if row is None:
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
kb_id = row["kb_id"]
thread_id = row["thread_id"]
if kb_id is not None:
owner_row = conn.execute(
"SELECT owner_user_id FROM rag_knowledge_bases WHERE id = ?",
(kb_id,),
).fetchone()
if owner_row is None:
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
owner = owner_row["owner_user_id"]
if owner is None or owner != current_subject:
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
return row
if thread_id is not None:
# Single-user invariant (see TODO above). We require the
# thread row to exist; an unknown thread_id is treated as
# not-found, not as silent grant.
thread_row = conn.execute(
"SELECT id FROM chat_threads WHERE id = ?",
(thread_id,),
).fetchone()
if thread_row is None:
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
return row
# Documents must belong to either a KB or a thread (DB CHECK
# constraint enforces XOR on insert); a row that satisfies
# neither is corrupt — treat as 404.
raise HTTPException(status_code=404, detail=_NOT_FOUND_DETAIL)
def chunk_belongs_to_document(chunk_id: str, document_id: str) -> bool:
"""True iff `chunk_id` exists in `rag_chunks` for `document_id`.
Used by `/preview-target?chunk_id=...` after the caller has
already established subject authorization for `document_id`. Does
NOT perform authorization itself: callers MUST call
`document_for_subject_or_404(document_id, ...)` first, otherwise a
valid `chunk_id` from another subject's document would leak via a
`True` return.
"""
if not chunk_id or not document_id:
return False
with get_connection() as conn:
row = conn.execute(
"SELECT 1 FROM rag_chunks WHERE id = ? AND document_id = ?",
(chunk_id, document_id),
).fetchone()
return row is not None

View file

@ -48,6 +48,11 @@ class Chunk:
text: str
token_count: int
page_number: int | None = None
source_page_index: int | None = None
page_char_start: int | None = None
page_char_end: int | None = None
line_start: int | None = None
line_end: int | None = None
TokenCounter = Callable[[str], int]
@ -132,6 +137,28 @@ def _merge(
return [c.strip() for c in chunks if c.strip()]
def _line_bounds(text: str, start: int, end: int) -> tuple[int, int]:
"""Return 1-based inclusive line numbers for a page-local span."""
line_start = text.count("\n", 0, start) + 1
line_end = text.count("\n", 0, max(start, end - 1)) + 1
return line_start, line_end
def _locate_piece(
page_text: str,
piece: str,
search_cursor: int,
) -> tuple[int | None, int | None, int | None, int | None, int]:
idx = page_text.find(piece, search_cursor)
if idx < 0:
idx = page_text.find(piece)
if idx < 0:
return None, None, None, None, search_cursor
end = idx + len(piece)
line_start, line_end = _line_bounds(page_text, idx, end)
return idx, end, line_start, line_end, idx + 1
# Markdown headings first so layout-aware parser output splits at sections.
DEFAULT_SEPARATORS: tuple[str, ...] = (
"\n# ",
@ -157,16 +184,27 @@ def chunk_pages(
"""Split pages independently so page_number stays attached to chunks."""
count = token_counter or _char_token_estimate
out: list[Chunk] = []
for page in pages:
for page_index, page in enumerate(pages):
search_cursor = 0
for segment in _split_at_figure_boundaries(page.text):
atomic = _atomic_split(segment, separators, max_tokens, count)
merged = _merge(atomic, max_tokens, overlap_tokens, count)
for piece in merged:
start, end, line_start, line_end, search_cursor = _locate_piece(
page.text,
piece,
search_cursor,
)
out.append(
Chunk(
text = piece,
token_count = count(piece),
page_number = page.page_number,
source_page_index = page_index,
page_char_start = start,
page_char_end = end,
line_start = line_start,
line_end = line_end,
)
)
return out
@ -192,13 +230,13 @@ def chunk_pages_with_spans(
count = token_counter or _char_token_estimate
parts: list[str] = []
page_ranges: list[tuple[int, int, int | None]] = []
page_ranges: list[tuple[int, int, int, int | None]] = []
cursor = 0
for index, page in enumerate(pages):
parts.append(page.text)
start = cursor
end = cursor + len(page.text)
page_ranges.append((start, end, page.page_number))
page_ranges.append((start, end, index, page.page_number))
cursor = end
if index < len(pages) - 1:
cursor += len(_PAGE_SEPARATOR)
@ -223,12 +261,34 @@ def chunk_pages_with_spans(
if idx < 0:
continue
end_idx = idx + len(text)
page_number = _page_for_span(idx, end_idx, page_ranges)
page_locator = _page_for_span(idx, end_idx, page_ranges)
source_page_index: int | None = None
page_number: int | None = None
page_char_start: int | None = None
page_char_end: int | None = None
line_start: int | None = None
line_end: int | None = None
if page_locator is not None:
page_start, page_end, page_idx, page_no = page_locator
source_page_index = page_idx
page_number = page_no
page_char_start = max(0, idx - page_start)
page_char_end = min(page_end, end_idx) - page_start
line_start, line_end = _line_bounds(
pages[page_idx].text,
page_char_start,
page_char_end,
)
chunks.append(
Chunk(
text = text,
token_count = count(text),
page_number = page_number,
source_page_index = source_page_index,
page_char_start = page_char_start,
page_char_end = page_char_end,
line_start = line_start,
line_end = line_end,
)
)
char_spans.append((idx, end_idx))
@ -241,9 +301,9 @@ def chunk_pages_with_spans(
def _page_for_span(
start: int,
end: int,
page_ranges: list[tuple[int, int, int | None]],
) -> int | None:
for ps, pe, pn in page_ranges:
page_ranges: list[tuple[int, int, int, int | None]],
) -> tuple[int, int, int, int | None] | None:
for ps, pe, page_index, page_number in page_ranges:
if start < pe and end > ps:
return pn
return ps, pe, page_index, page_number
return None

View file

@ -9,6 +9,7 @@ vectors, and rebuilds BM25 on completion. Only the parent opens rag.db.
from __future__ import annotations
import json
import multiprocessing as mp
import queue as queue_module
import sqlite3
@ -97,6 +98,22 @@ def _subprocess_worker(
)
pages = inline_image_captions(pages, parsed.images, captions)
out_queue.put(
{
"type": "document_pages",
"pages": [
{
"page_index": index,
"page_number": page.page_number,
"text": page.text,
"char_count": len(page.text),
"line_count": len(page.text.splitlines()),
}
for index, page in enumerate(pages)
],
}
)
out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1})
from core.rag.embeddings import (
get_embedder,
@ -112,6 +129,7 @@ def _subprocess_worker(
if chunking_strategy == "late":
_run_late_chunking(
pages = pages,
stored_path = Path(stored_path),
chunk_size = chunk_size,
overlap = overlap,
counter = counter,
@ -123,6 +141,7 @@ def _subprocess_worker(
text_count = _run_standard_chunking(
pages = pages,
stored_path = Path(stored_path),
chunk_size = chunk_size,
overlap = overlap,
counter = counter,
@ -151,6 +170,7 @@ def _subprocess_worker(
def _run_standard_chunking(
*,
pages,
stored_path,
chunk_size,
overlap,
counter,
@ -172,6 +192,9 @@ def _run_standard_chunking(
if send_complete:
out_queue.put({"type": "error", "error": "chunker produced no chunks"})
return 0
from core.rag.locators import pdf_regions_for_chunks
pdf_regions = pdf_regions_for_chunks(stored_path, pages, chunks)
total = len(chunks)
for i in range(0, total, batch_size):
@ -192,9 +215,15 @@ def _run_standard_chunking(
"text": c.text,
"token_count": c.token_count,
"page_number": c.page_number,
"source_page_index": c.source_page_index,
"page_char_start": c.page_char_start,
"page_char_end": c.page_char_end,
"line_start": c.line_start,
"line_end": c.line_end,
"pdf_regions": pdf_regions[i + offset],
"kind": "text",
}
for c in batch
for offset, c in enumerate(batch)
],
"vectors": vectors.tolist(),
}
@ -315,6 +344,7 @@ def _stream_image_chunks(
def _run_late_chunking(
*,
pages,
stored_path,
chunk_size,
overlap,
counter,
@ -324,6 +354,7 @@ def _run_late_chunking(
) -> None:
"""Chunk once, embed in one pass, ship all chunks in one chunks_batch."""
from core.rag.chunking import chunk_pages_with_spans
from core.rag.locators import pdf_regions_for_chunks
out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2})
full_doc, chunks, char_spans = chunk_pages_with_spans(
@ -343,6 +374,7 @@ def _run_late_chunking(
model_name = model_name,
normalize = True,
)
pdf_regions = pdf_regions_for_chunks(stored_path, pages, chunks)
out_queue.put({"type": "progress", "stage": "embed", "progress": 0.9})
out_queue.put(
@ -354,8 +386,15 @@ def _run_late_chunking(
"text": c.text,
"token_count": c.token_count,
"page_number": c.page_number,
"source_page_index": c.source_page_index,
"page_char_start": c.page_char_start,
"page_char_end": c.page_char_end,
"line_start": c.line_start,
"line_end": c.line_end,
"pdf_regions": pdf_regions[index],
"kind": "text",
}
for c in chunks
for index, c in enumerate(chunks)
],
"vectors": [v.tolist() for v in vectors],
}
@ -469,6 +508,14 @@ def _insert_chunks_and_collect_for_bm25(
meta["page_number"],
kind,
image_path,
meta.get("source_page_index"),
meta.get("page_char_start"),
meta.get("page_char_end"),
meta.get("line_start"),
meta.get("line_end"),
json.dumps(meta.get("pdf_regions") or [], separators = (",", ":"))
if meta.get("pdf_regions")
else None,
)
)
points.append(
@ -482,6 +529,12 @@ def _insert_chunks_and_collect_for_bm25(
"page_number": meta["page_number"],
"kind": kind,
"image_path": image_path,
"source_page_index": meta.get("source_page_index"),
"page_char_start": meta.get("page_char_start"),
"page_char_end": meta.get("page_char_end"),
"line_start": meta.get("line_start"),
"line_end": meta.get("line_end"),
"pdf_regions": meta.get("pdf_regions") or [],
},
}
)
@ -492,8 +545,9 @@ def _insert_chunks_and_collect_for_bm25(
"""
INSERT INTO rag_chunks
(id, document_id, chunk_index, text, token_count, page_number,
kind, image_path)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
kind, image_path, source_page_index, page_char_start,
page_char_end, line_start, line_end, pdf_regions_json)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
rows,
)
@ -515,6 +569,41 @@ def _insert_chunks_and_collect_for_bm25(
return bm25_rows
def _replace_document_pages(document_id: str, pages: list[dict]) -> None:
now = int(time.time())
rows = [
(
document_id,
int(page["page_index"]),
page.get("page_number"),
page.get("text") or "",
int(page.get("char_count", len(page.get("text") or ""))),
int(page.get("line_count", len((page.get("text") or "").splitlines()))),
now,
)
for page in pages
]
with get_connection() as conn:
doc_row = conn.execute(
"SELECT 1 FROM rag_documents WHERE id = ?",
(document_id,),
).fetchone()
if doc_row is None:
raise sqlite3.IntegrityError("FOREIGN KEY constraint failed")
conn.execute("DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,))
if rows:
conn.executemany(
"""
INSERT INTO rag_document_pages
(document_id, page_index, page_number, text, char_count,
line_count, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
rows,
)
conn.commit()
def _all_scope_chunks(scope: str) -> list[dict]:
if scope.startswith("kb_"):
kb_id = scope[len("kb_") :]
@ -579,6 +668,17 @@ def _pump(
elif mtype == "dim":
embedding_dim = int(msg["dim"])
vector_store.ensure_collection(state.scope, embedding_dim)
elif mtype == "document_pages":
try:
_replace_document_pages(
state.document_id,
list(msg.get("pages") or []),
)
except sqlite3.IntegrityError as exc:
final_error = (
f"document was removed before ingestion finished ({exc})"
)
break
elif mtype == "chunks_batch":
if embedding_dim is None:
embedding_dim = len(msg["vectors"][0]) if msg["vectors"] else None

View file

@ -0,0 +1,494 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Backfill and PDF-region helpers for durable RAG chunk locators."""
from __future__ import annotations
import json
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from loggers import get_logger
from storage.studio_db import get_connection
from . import vector_store
from .parsers import ParsedPage, parse
from .vector_store import kb_scope, thread_scope
logger = get_logger(__name__)
@dataclass(frozen = True)
class LocatorMatch:
page_index: int
page_number: int | None
start: int
end: int
line_start: int
line_end: int
@dataclass(frozen = True)
class BackfillResult:
document_id: str
total_chunks: int
matched: int
already_located: int
ambiguous: int
missing: int
skipped: int
regions_matched: int
pages_refreshed: int
def _line_bounds(text: str, start: int, end: int) -> tuple[int, int]:
line_start = text.count("\n", 0, start) + 1
line_end = text.count("\n", 0, max(start, end - 1)) + 1
return line_start, line_end
def _find_exact(page_text: str, needle: str) -> list[tuple[int, int]]:
if not needle:
return []
out: list[tuple[int, int]] = []
cursor = 0
while True:
idx = page_text.find(needle, cursor)
if idx < 0:
break
out.append((idx, idx + len(needle)))
cursor = idx + 1
return out
def _normalize_with_map(text: str) -> tuple[str, list[int], list[int]]:
chars: list[str] = []
starts: list[int] = []
ends: list[int] = []
last_space = False
for idx, ch in enumerate(text):
if ch.isspace():
if chars and not last_space:
chars.append(" ")
starts.append(idx)
ends.append(idx + 1)
elif chars and last_space:
ends[-1] = idx + 1
last_space = True
continue
chars.append(ch.casefold())
starts.append(idx)
ends.append(idx + 1)
last_space = False
first = 0
while first < len(chars) and chars[first] == " ":
first += 1
last = len(chars)
while last > first and chars[last - 1] == " ":
last -= 1
return "".join(chars[first:last]), starts[first:last], ends[first:last]
def _find_normalized(page_text: str, needle: str) -> list[tuple[int, int]]:
norm_page, starts, ends = _normalize_with_map(page_text)
norm_needle, _needle_starts, _needle_ends = _normalize_with_map(needle)
if not norm_page or not norm_needle:
return []
out: list[tuple[int, int]] = []
cursor = 0
while True:
idx = norm_page.find(norm_needle, cursor)
if idx < 0:
break
end_idx = idx + len(norm_needle) - 1
if 0 <= idx < len(starts) and 0 <= end_idx < len(ends):
out.append((starts[idx], ends[end_idx]))
cursor = idx + 1
return out
def _locate_unique(text: str, pages: list[ParsedPage]) -> tuple[LocatorMatch | None, str]:
text = (text or "").strip()
if not text:
return None, "missing"
matches: list[LocatorMatch] = []
for page_index, page in enumerate(pages):
for start, end in _find_exact(page.text, text):
line_start, line_end = _line_bounds(page.text, start, end)
matches.append(
LocatorMatch(
page_index = page_index,
page_number = page.page_number,
start = start,
end = end,
line_start = line_start,
line_end = line_end,
)
)
if len(matches) == 1:
return matches[0], "matched"
if len(matches) > 1:
return None, "ambiguous"
for page_index, page in enumerate(pages):
for start, end in _find_normalized(page.text, text):
line_start, line_end = _line_bounds(page.text, start, end)
matches.append(
LocatorMatch(
page_index = page_index,
page_number = page.page_number,
start = start,
end = end,
line_start = line_start,
line_end = line_end,
)
)
if len(matches) == 1:
return matches[0], "matched"
if len(matches) > 1:
return None, "ambiguous"
return None, "missing"
def _replace_document_pages(document_id: str, pages: list[ParsedPage]) -> None:
now = int(time.time())
rows = [
(
document_id,
index,
page.page_number,
page.text,
len(page.text),
len(page.text.splitlines()),
now,
)
for index, page in enumerate(pages)
]
with get_connection() as conn:
conn.execute("DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,))
if rows:
conn.executemany(
"""
INSERT INTO rag_document_pages
(document_id, page_index, page_number, text, char_count,
line_count, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
rows,
)
conn.commit()
def _region_anchor(page_text: str, match: LocatorMatch) -> str | None:
segment = page_text[match.start : match.end]
words = [w.strip(" \t\r\n*#`[]()") for w in segment.split()]
words = [w for w in words if len(w) >= 2]
if len(words) < 3:
return None
anchor = " ".join(words[: min(16, len(words))])
return anchor if len(anchor) >= 12 else None
def _normalized_occurrences(haystack: str, needle: str) -> int:
norm_haystack, _starts, _ends = _normalize_with_map(haystack)
norm_needle, _needle_starts, _needle_ends = _normalize_with_map(needle)
if not norm_haystack or not norm_needle:
return 0
count = 0
cursor = 0
while True:
idx = norm_haystack.find(norm_needle, cursor)
if idx < 0:
return count
count += 1
cursor = idx + 1
def pdf_regions_for_match(
pdf_path: Path,
pages: list[ParsedPage],
match: LocatorMatch,
) -> list[dict[str, Any]]:
"""Return normalized PDF rectangles for a unique chunk match.
Regions are intentionally conservative: no PyMuPDF, no page, no
unique anchor, or no positive-area rectangles all produce an empty
list rather than guessed highlights.
"""
if pdf_path.suffix.lower() != ".pdf":
return []
if match.page_index < 0 or match.page_index >= len(pages):
return []
anchor = _region_anchor(pages[match.page_index].text, match)
if not anchor:
return []
try:
import pymupdf
except Exception:
return []
try:
doc = pymupdf.open(str(pdf_path))
except Exception:
return []
try:
return _pdf_regions_for_match_doc(doc, pages, match, anchor)
finally:
doc.close()
def _pdf_regions_for_match_doc(
doc: Any,
pages: list[ParsedPage],
match: LocatorMatch,
anchor: str,
) -> list[dict[str, Any]]:
try:
if match.page_index >= len(doc):
return []
page = doc[match.page_index]
raw_text = page.get_text("text") or ""
if _normalized_occurrences(raw_text, anchor) != 1:
return []
rects = page.search_for(anchor) or []
page_rect = page.rect
page_width = float(page_rect.width)
page_height = float(page_rect.height)
if page_width <= 0 or page_height <= 0:
return []
out: list[dict[str, Any]] = []
for rect in rects:
width = max(0.0, float(rect.x1 - rect.x0))
height = max(0.0, float(rect.y1 - rect.y0))
if width <= 0 or height <= 0:
continue
out.append(
{
"pageIndex": match.page_index,
"pageNumber": match.page_number,
"x": max(0.0, min(1.0, float(rect.x0) / page_width)),
"y": max(0.0, min(1.0, float(rect.y0) / page_height)),
"width": max(0.0, min(1.0, width / page_width)),
"height": max(0.0, min(1.0, height / page_height)),
"confidence": "exact",
"source": "pymupdf-search",
}
)
return out
except Exception:
return []
def pdf_regions_for_chunks(
pdf_path: Path,
pages: list[ParsedPage],
chunks: list[Any],
) -> list[list[dict[str, Any]]]:
if pdf_path.suffix.lower() != ".pdf":
return [[] for _ in chunks]
try:
import pymupdf
doc = pymupdf.open(str(pdf_path))
except Exception:
return [[] for _ in chunks]
regions: list[list[dict[str, Any]]] = []
try:
for chunk in chunks:
page_index = getattr(chunk, "source_page_index", None)
start = getattr(chunk, "page_char_start", None)
end = getattr(chunk, "page_char_end", None)
if page_index is None or start is None or end is None:
regions.append([])
continue
if page_index < 0 or page_index >= len(pages):
regions.append([])
continue
line_start, line_end = _line_bounds(pages[page_index].text, start, end)
match = LocatorMatch(
page_index = int(page_index),
page_number = getattr(chunk, "page_number", None),
start = int(start),
end = int(end),
line_start = line_start,
line_end = line_end,
)
anchor = _region_anchor(pages[match.page_index].text, match)
if not anchor:
regions.append([])
continue
regions.append(_pdf_regions_for_match_doc(doc, pages, match, anchor))
return regions
finally:
doc.close()
def _scope_for_document(kb_id: str | None, thread_id: str | None) -> str | None:
if kb_id:
return kb_scope(kb_id)
if thread_id:
return thread_scope(thread_id)
return None
def _update_vector_payloads(scope: str | None, updates: dict[str, dict[str, Any]]) -> None:
if not scope or not updates:
return
try:
vector_store.update_chunk_payload_fields(scope, updates)
except Exception as exc:
logger.warning(
"RAG locator backfill: vector payload update failed",
error = str(exc),
)
def backfill_document_locators(document_id: str, stored_path: Path) -> BackfillResult:
parsed = parse(stored_path, want_images = False)
pages = parsed.pages
_replace_document_pages(document_id, pages)
with get_connection() as conn:
doc_row = conn.execute(
"SELECT kb_id, thread_id FROM rag_documents WHERE id = ?",
(document_id,),
).fetchone()
if doc_row is None:
return BackfillResult(document_id, 0, 0, 0, 0, 0, 0, 0, len(pages))
rows = conn.execute(
"""
SELECT id, text, kind, page_number, source_page_index,
page_char_start, page_char_end, line_start, line_end,
pdf_regions_json
FROM rag_chunks
WHERE document_id = ?
ORDER BY chunk_index ASC
""",
(document_id,),
).fetchall()
scope = _scope_for_document(doc_row["kb_id"], doc_row["thread_id"])
total = len(rows)
matched = 0
already_located = 0
ambiguous = 0
missing = 0
skipped = 0
regions_matched = 0
sql_updates: list[tuple[Any, ...]] = []
vector_updates: dict[str, dict[str, Any]] = {}
for row in rows:
kind = row["kind"] or "text"
text = row["text"] or ""
if kind not in ("text", "caption") or not text.strip():
skipped += 1
continue
existing_complete = (
row["source_page_index"] is not None
and row["page_char_start"] is not None
and row["page_char_end"] is not None
and row["line_start"] is not None
and row["line_end"] is not None
)
match: LocatorMatch | None
status: str
if existing_complete:
already_located += 1
page_index = int(row["source_page_index"])
if 0 <= page_index < len(pages):
match = LocatorMatch(
page_index = page_index,
page_number = row["page_number"],
start = int(row["page_char_start"]),
end = int(row["page_char_end"]),
line_start = int(row["line_start"]),
line_end = int(row["line_end"]),
)
else:
match = None
status = "already_located"
else:
match, status = _locate_unique(text, pages)
if status == "matched" and match is not None:
matched += 1
elif status == "ambiguous":
ambiguous += 1
continue
else:
missing += 1
continue
if match is None:
continue
regions = pdf_regions_for_match(stored_path, pages, match)
regions_json = json.dumps(regions, separators = (",", ":")) if regions else None
if regions:
regions_matched += 1
if status == "matched" or (regions and not row["pdf_regions_json"]):
sql_updates.append(
(
match.page_number,
match.page_index,
match.start,
match.end,
match.line_start,
match.line_end,
regions_json,
row["id"],
)
)
vector_updates[row["id"]] = {
"page_number": match.page_number,
"source_page_index": match.page_index,
"page_char_start": match.start,
"page_char_end": match.end,
"line_start": match.line_start,
"line_end": match.line_end,
"pdf_regions": regions,
}
if sql_updates:
with get_connection() as conn:
conn.executemany(
"""
UPDATE rag_chunks
SET page_number = COALESCE(page_number, ?),
source_page_index = ?,
page_char_start = ?,
page_char_end = ?,
line_start = ?,
line_end = ?,
pdf_regions_json = COALESCE(?, pdf_regions_json)
WHERE id = ?
""",
sql_updates,
)
conn.commit()
_update_vector_payloads(scope, vector_updates)
return BackfillResult(
document_id = document_id,
total_chunks = total,
matched = matched,
already_located = already_located,
ambiguous = ambiguous,
missing = missing,
skipped = skipped,
regions_matched = regions_matched,
pages_refreshed = len(pages),
)

View file

@ -50,6 +50,11 @@ class Hit:
document_id: str | None = None
chunk_index: int | None = None
kind: str = "text"
source_page_index: int | None = None
page_char_start: int | None = None
page_char_end: int | None = None
line_start: int | None = None
line_end: int | None = None
# Raw cosine; None for BM25-only hits.
dense_score: float | None = None
@ -151,6 +156,11 @@ def retrieve_dense(
document_id = payload.get("document_id"),
chunk_index = payload.get("chunk_index"),
kind = payload.get("kind", "text"),
source_page_index = payload.get("source_page_index"),
page_char_start = payload.get("page_char_start"),
page_char_end = payload.get("page_char_end"),
line_start = payload.get("line_start"),
line_end = payload.get("line_end"),
dense_score = r["score"],
)
)
@ -184,6 +194,11 @@ def _rrf_fuse(
document_id = seen[cid].document_id,
chunk_index = seen[cid].chunk_index,
kind = seen[cid].kind,
source_page_index = seen[cid].source_page_index,
page_char_start = seen[cid].page_char_start,
page_char_end = seen[cid].page_char_end,
line_start = seen[cid].line_start,
line_end = seen[cid].line_end,
dense_score = dense_scores.get(cid),
)
for cid, score in ordered

View file

@ -93,12 +93,33 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
f'id="{index}"',
f'source="{_xml_attr(hit.get("filename") or "unknown")}"',
]
# Durable backend ids — additive per contracts.md §3.1 (T3).
# ``id`` above stays as the visible citation id (used by the model
# as `[N]`); ``document_id`` + ``chunk_id`` are what the preview
# route consumes. Old XML without these attrs still parses on
# the frontend (hover-only), per contracts §3.2.
document_id = hit.get("document_id")
if document_id:
attrs.append(f'document_id="{_xml_attr(document_id)}"')
backend_chunk_id = hit.get("chunk_id")
if backend_chunk_id:
attrs.append(f'chunk_id="{_xml_attr(backend_chunk_id)}"')
page = hit.get("page_number")
if page is not None:
attrs.append(f'page="{page}"')
chunk_index = hit.get("chunk_index")
if chunk_index is not None:
attrs.append(f'chunk_index="{chunk_index}"')
for attr_name in (
"source_page_index",
"page_char_start",
"page_char_end",
"line_start",
"line_end",
):
value = hit.get(attr_name)
if value is not None:
attrs.append(f'{attr_name}="{value}"')
tokens = hit.get("token_count")
if tokens:
attrs.append(f'tokens="{tokens}"')
@ -106,7 +127,6 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
if kind and kind != "text":
attrs.append(f'kind="{_xml_attr(kind)}"')
image_path = hit.get("image_path")
document_id = hit.get("document_id")
if kind == "image" and image_path and document_id:
# Mirror routes/rag.py search-response shape so the frontend
# tool card can render the image inline via the same route.
@ -213,6 +233,8 @@ def search_knowledge_base(
f"""
SELECT c.id AS chunk_id, c.text, c.page_number,
c.token_count, c.kind, c.image_path,
c.source_page_index, c.page_char_start,
c.page_char_end, c.line_start, c.line_end,
c.document_id, d.filename
FROM rag_chunks c
JOIN rag_documents d ON d.id = c.document_id

View file

@ -131,6 +131,45 @@ def search(
return out
def update_chunk_payload_fields(
scope: str,
updates: dict[str, dict],
) -> None:
"""Merge locator fields into existing vector payload JSON by chunk id."""
from core.rag.db import get_rag_connection
if not updates:
return
conn = get_rag_connection()
rows = conn.execute(
f"""
SELECT chunk_id, payload_json
FROM rag_vectors
WHERE scope = ? AND chunk_id IN ({",".join("?" for _ in updates)})
""",
[scope, *updates.keys()],
).fetchall()
payload_rows: list[tuple[str, str]] = []
for row in rows:
try:
payload = json.loads(row["payload_json"] or "{}")
except json.JSONDecodeError:
payload = {}
payload.update(updates.get(row["chunk_id"], {}))
payload_rows.append((json.dumps(payload, default = str), row["chunk_id"], scope))
if not payload_rows:
return
conn.executemany(
"""
UPDATE rag_vectors
SET payload_json = ?
WHERE chunk_id = ? AND scope = ?
""",
payload_rows,
)
conn.commit()
def delete_scope(scope: str) -> None:
from core.rag.db import get_rag_connection

View file

@ -10,10 +10,13 @@ import json
import os
import queue as queue_module
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Literal, Optional
from urllib.parse import quote
from uuid import uuid4
import jwt
from fastapi import (
APIRouter,
Depends,
@ -23,10 +26,11 @@ from fastapi import (
Request,
UploadFile,
)
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.responses import FileResponse, Response, StreamingResponse
from pydantic import BaseModel, Field
from auth.authentication import get_current_subject, get_current_subject_sse
from auth.storage import get_jwt_secret
async def _sse_auth(
@ -37,6 +41,8 @@ async def _sse_auth(
from core.rag import embeddings, ingestion, reranker, retrieval, vector_store
from core.rag.authorization import document_for_subject_or_404
from core.rag.locators import backfill_document_locators
from core.rag.vector_store import kb_scope, thread_scope
from loggers import get_logger
from storage.studio_db import (
@ -44,7 +50,7 @@ from storage.studio_db import (
list_chat_settings,
upsert_chat_settings_merge,
)
from utils.paths.storage_roots import ensure_dir, rag_uploads_root
from utils.paths.storage_roots import ensure_dir, rag_uploads_root, resolve_under_root
from utils.rag.config import (
RAG_MAX_UPLOAD_MB,
RAG_RERANK_CANDIDATE_K,
@ -139,6 +145,11 @@ class SearchHit(BaseModel):
filename: str | None = None
kind: str = "text"
image_url: str | None = None
source_page_index: int | None = None
page_char_start: int | None = None
page_char_end: int | None = None
line_start: int | None = None
line_end: int | None = None
class SearchResponse(BaseModel):
@ -872,6 +883,7 @@ def get_rag_image(
current_subject: str = Depends(get_current_subject),
) -> FileResponse:
"""Serve an extracted image; realpath-check against the uploads root."""
document_for_subject_or_404(document_id, current_subject)
if "/" in filename or "\\" in filename or filename.startswith("."):
raise HTTPException(status_code = 400, detail = "Invalid filename")
root = Path(os.path.realpath(rag_uploads_root() / "images"))
@ -1016,6 +1028,538 @@ async def _replay_terminal_state(row: Any):
# --- Search ---
# --- Document preview (file + preview-target) ---
PreviewMediaKind = Literal["pdf", "text", "docx", "html", "image", "unknown"]
PreviewChunkKind = Literal["text", "image", "caption"]
class PreviewPdfRegion(BaseModel):
pageIndex: int
pageNumber: int | None = None
x: float
y: float
width: float
height: float
confidence: Literal["exact"]
source: str
class PreviewTargetResponse(BaseModel):
"""Per contracts.md §1.2 / §1.3 — single shape covering both
cited-chunk and document-row preview modes. The §1.3 metadata-only
mode returns ``None`` for ``chunkId``/``chunkIndex``/``targetPage``/
``snippet``/``kind``/``imageUrl`` (Q2: no first-chunk guessing).
"""
documentId: str
filename: str
contentType: str | None
mediaKind: PreviewMediaKind
byteSize: int
status: str
kbId: str | None
threadId: str | None
chunkId: str | None
chunkIndex: int | None
targetPage: int | None
snippet: str | None
kind: PreviewChunkKind | None
imageUrl: str | None
sourcePageIndex: int | None
pageCharStart: int | None
pageCharEnd: int | None
lineStart: int | None
lineEnd: int | None
pdfRegions: list[PreviewPdfRegion] = Field(default_factory = list)
class PreviewFileUrlResponse(BaseModel):
url: str
expiresAt: int
class LocatorBackfillResponse(BaseModel):
documentId: str
totalChunks: int
matched: int
alreadyLocated: int
ambiguous: int
missing: int
skipped: int
regionsMatched: int
pagesRefreshed: int
# Extension allowlist for inline rendering / disposition. Anything not in
# this map collapses to ("application/octet-stream", attachment, "unknown").
# .html / .htm intentionally serve as text/plain attachment (decisions Q7 +
# Risk #3) so an uploaded HTML cannot execute in the app origin.
_PREVIEW_EXT_MAP: dict[str, tuple[str, str, PreviewMediaKind]] = {
".pdf": ("application/pdf", "inline", "pdf"),
".txt": ("text/plain; charset=utf-8", "inline", "text"),
".md": ("text/markdown; charset=utf-8", "inline", "text"),
".markdown": ("text/markdown; charset=utf-8", "inline", "text"),
".docx": (
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
"attachment",
"docx",
),
".html": ("text/plain; charset=utf-8", "attachment", "html"),
".htm": ("text/plain; charset=utf-8", "attachment", "html"),
".png": ("image/png", "inline", "image"),
".jpg": ("image/jpeg", "inline", "image"),
".jpeg": ("image/jpeg", "inline", "image"),
".gif": ("image/gif", "inline", "image"),
".webp": ("image/webp", "inline", "image"),
}
def _ascii_only(value: str) -> bool:
try:
value.encode("ascii")
except UnicodeEncodeError:
return False
return True
def _content_disposition_header(filename: str, disposition: str) -> str:
"""Build a Content-Disposition header. Non-ASCII filenames use RFC 5987
``filename*=UTF-8''`` alongside an ASCII-only ``filename=`` fallback so
older clients still get something readable.
"""
from urllib.parse import quote as _urlquote
safe = filename.replace('"', "").replace("\r", "").replace("\n", "")
if _ascii_only(safe):
return f'{disposition}; filename="{safe}"'
ascii_fallback = safe.encode("ascii", "replace").decode("ascii")
encoded = _urlquote(safe, safe = "")
return (
f'{disposition}; filename="{ascii_fallback}"; '
f"filename*=UTF-8''{encoded}"
)
def _preview_file_metadata(filename: str) -> tuple[str, str, PreviewMediaKind]:
"""Map a stored filename to (content_type, disposition, mediaKind).
Unknown extensions always force ``application/octet-stream`` +
``attachment`` + ``unknown`` so the browser cannot sniff a sensitive
type and inline it (Risk #3).
"""
ext = Path(filename).suffix.lower()
return _PREVIEW_EXT_MAP.get(
ext, ("application/octet-stream", "attachment", "unknown")
)
_PREVIEW_FILE_AUDIENCE = "rag-preview-file"
_PREVIEW_FILE_TTL_SECONDS = 5 * 60
_JWT_ALGORITHM = "HS256"
def _parse_pdf_regions(value: str | None) -> list[PreviewPdfRegion]:
if not value:
return []
try:
raw = json.loads(value)
except (TypeError, json.JSONDecodeError):
return []
if not isinstance(raw, list):
return []
out: list[PreviewPdfRegion] = []
for item in raw:
if not isinstance(item, dict):
continue
try:
region = PreviewPdfRegion(**item)
except Exception:
continue
if (
0 <= region.x <= 1
and 0 <= region.y <= 1
and region.width > 0
and region.height > 0
):
out.append(region)
return out
def _preview_file_token(
*,
subject: str,
document_id: str,
) -> tuple[str, int]:
secret = get_jwt_secret(subject)
if secret is None:
raise HTTPException(status_code = 401, detail = "Invalid or expired token")
expires = datetime.now(timezone.utc) + timedelta(seconds = _PREVIEW_FILE_TTL_SECONDS)
payload = {
"sub": subject,
"aud": _PREVIEW_FILE_AUDIENCE,
"document_id": document_id,
"exp": expires,
}
token = jwt.encode(payload, secret, algorithm = _JWT_ALGORITHM)
return token, int(expires.timestamp())
def _subject_from_preview_file_token(document_id: str, token: str) -> str:
try:
unverified = jwt.decode(
token,
options = {
"verify_signature": False,
"verify_exp": False,
"verify_aud": False,
},
)
except jwt.InvalidTokenError as exc:
raise HTTPException(status_code = 401, detail = "Invalid preview token") from exc
subject = unverified.get("sub")
if not isinstance(subject, str) or not subject:
raise HTTPException(status_code = 401, detail = "Invalid preview token")
secret = get_jwt_secret(subject)
if secret is None:
raise HTTPException(status_code = 401, detail = "Invalid preview token")
try:
payload = jwt.decode(
token,
secret,
algorithms = [_JWT_ALGORITHM],
audience = _PREVIEW_FILE_AUDIENCE,
)
except jwt.InvalidTokenError as exc:
raise HTTPException(status_code = 401, detail = "Invalid preview token") from exc
if payload.get("document_id") != document_id:
raise HTTPException(status_code = 401, detail = "Invalid preview token")
return subject
def _resolve_document_file_or_404(doc_row: Any, document_id: str) -> Path:
try:
resolved = resolve_under_root(
doc_row["stored_path"],
root = rag_uploads_root(),
)
except ValueError as exc:
# Symlink escape / ``..`` / absolute outside root — collapse to
# "file not found" (the auth row exists, the bytes do not).
logger.warning(
"RAG preview: stored_path escaped uploads root for doc %s: %s",
document_id,
exc,
)
raise HTTPException(
status_code = 404,
detail = "Document file not found",
) from exc
if not resolved.is_file():
raise HTTPException(
status_code = 404,
detail = "Document file not found",
)
return resolved
def _parse_range_header(range_header: str | None, size: int) -> tuple[int, int] | None:
if not range_header:
return None
if not range_header.startswith("bytes="):
raise ValueError("unsupported range unit")
spec = range_header[len("bytes=") :].strip()
if "," in spec or "-" not in spec:
raise ValueError("multiple or malformed ranges are not supported")
start_s, end_s = spec.split("-", 1)
if not start_s and not end_s:
raise ValueError("empty range")
if not start_s:
suffix = int(end_s)
if suffix <= 0:
raise ValueError("invalid suffix range")
start = max(0, size - suffix)
end = size - 1
else:
start = int(start_s)
end = int(end_s) if end_s else size - 1
if start < 0 or end < start or start >= size:
raise ValueError("range outside file")
return start, min(end, size - 1)
def _iter_file_range(path: Path, start: int, end: int):
with path.open("rb") as fh:
fh.seek(start)
remaining = end - start + 1
while remaining > 0:
chunk = fh.read(min(64 * 1024, remaining))
if not chunk:
break
remaining -= len(chunk)
yield chunk
def _serve_document_file_row(
doc_row: Any,
document_id: str,
range_header: str | None,
) -> FileResponse | Response | StreamingResponse:
resolved = _resolve_document_file_or_404(doc_row, document_id)
content_type, disposition, _media_kind = _preview_file_metadata(
doc_row["filename"]
)
safe_name = _sanitize_filename(doc_row["filename"])
headers = {
"Content-Disposition": _content_disposition_header(safe_name, disposition),
"X-Content-Type-Options": "nosniff",
"Cache-Control": "private, max-age=0, must-revalidate",
"Accept-Ranges": "bytes",
}
size = resolved.stat().st_size
try:
byte_range = _parse_range_header(range_header, size)
except (TypeError, ValueError):
range_headers = dict(headers)
range_headers["Content-Range"] = f"bytes */{size}"
return Response(status_code = 416, headers = range_headers)
if byte_range is not None:
start, end = byte_range
range_headers = dict(headers)
range_headers["Content-Range"] = f"bytes {start}-{end}/{size}"
range_headers["Content-Length"] = str(end - start + 1)
return StreamingResponse(
_iter_file_range(resolved, start, end),
status_code = 206,
media_type = content_type,
headers = range_headers,
)
# Starlette's FileResponse handles ordinary downloads efficiently. We
# still advertise Accept-Ranges so PDF.js can switch to explicit range
# requests via the signed URL path.
return FileResponse(
path = str(resolved),
media_type = content_type,
headers = headers,
)
@router.get(
"/documents/{document_id}/preview-target",
response_model = PreviewTargetResponse,
)
def get_document_preview_target(
document_id: str,
chunk_id: Optional[str] = Query(None),
current_subject: str = Depends(get_current_subject),
) -> PreviewTargetResponse:
"""Resolve preview metadata for a document, optionally focused on a chunk.
See contracts.md §1 for the response shape. ``chunk_id`` is the durable
``rag_chunks.id`` carried as ``backendChunkId`` on the frontend; when
supplied it MUST belong to ``document_id`` or we collapse to 404 (so a
probe cannot enumerate cross-document chunk ids).
"""
doc_row = document_for_subject_or_404(document_id, current_subject)
_ct, _disposition, media_kind = _preview_file_metadata(doc_row["filename"])
base = {
"documentId": doc_row["id"],
"filename": doc_row["filename"],
"contentType": doc_row["content_type"],
"mediaKind": media_kind,
"byteSize": int(doc_row["byte_size"]),
"status": doc_row["status"],
"kbId": doc_row["kb_id"],
"threadId": doc_row["thread_id"],
}
if not chunk_id:
# Q2: document-row preview returns metadata only — frontend MUST
# NOT fall back to "first chunk".
return PreviewTargetResponse(
**base,
chunkId = None,
chunkIndex = None,
targetPage = None,
snippet = None,
kind = None,
imageUrl = None,
sourcePageIndex = None,
pageCharStart = None,
pageCharEnd = None,
lineStart = None,
lineEnd = None,
pdfRegions = [],
)
# Single connection enforces membership AND fetches the row in one
# query. Splitting this into a separate `chunk_belongs_to_document`
# call would open a second SQLite connection and create a TOCTOU
# window — if the chunk is deleted between the two calls, the data
# fetch returns None and the route 500s on the next attribute access
# (devils-advocate D1.1). The cross-document case still collapses to
# the same 404 the auth helper emits — never 400 (would leak doc
# existence).
with get_connection() as conn:
chunk_row = conn.execute(
"""
SELECT id, chunk_index, page_number, text, kind, image_path,
source_page_index, page_char_start, page_char_end,
line_start, line_end, pdf_regions_json
FROM rag_chunks WHERE id = ? AND document_id = ?
""",
(chunk_id, document_id),
).fetchone()
if chunk_row is None:
raise HTTPException(
status_code = 404,
detail = "Document not found",
)
chunk_kind: PreviewChunkKind = (chunk_row["kind"] or "text") # type: ignore[assignment]
image_url: str | None = None
if chunk_kind == "image" and chunk_row["image_path"]:
image_url = (
f"/api/rag/images/{doc_row['id']}/"
f"{Path(chunk_row['image_path']).name}"
)
return PreviewTargetResponse(
**base,
chunkId = chunk_row["id"],
chunkIndex = int(chunk_row["chunk_index"]),
targetPage = (
int(chunk_row["page_number"])
if chunk_row["page_number"] is not None
else None
),
snippet = chunk_row["text"] or "",
kind = chunk_kind,
imageUrl = image_url,
sourcePageIndex = chunk_row["source_page_index"],
pageCharStart = chunk_row["page_char_start"],
pageCharEnd = chunk_row["page_char_end"],
lineStart = chunk_row["line_start"],
lineEnd = chunk_row["line_end"],
pdfRegions = _parse_pdf_regions(chunk_row["pdf_regions_json"]),
)
@router.post(
"/documents/{document_id}/locators/backfill",
response_model = LocatorBackfillResponse,
)
def backfill_document_locators_route(
document_id: str,
current_subject: str = Depends(get_current_subject),
) -> LocatorBackfillResponse:
"""In-place locator backfill for existing citations.
This preserves ``document_id`` and ``chunk_id``. Chunks are updated only
when their text has one unambiguous match in the parsed document text;
duplicate or missing matches remain null.
"""
doc_row = document_for_subject_or_404(document_id, current_subject)
resolved = _resolve_document_file_or_404(doc_row, document_id)
try:
result = backfill_document_locators(document_id, resolved)
except Exception as exc:
logger.warning(
"RAG locator backfill failed",
document_id = document_id,
error = str(exc),
)
raise HTTPException(
status_code = 400,
detail = "Document locators could not be backfilled",
) from exc
return LocatorBackfillResponse(
documentId = result.document_id,
totalChunks = result.total_chunks,
matched = result.matched,
alreadyLocated = result.already_located,
ambiguous = result.ambiguous,
missing = result.missing,
skipped = result.skipped,
regionsMatched = result.regions_matched,
pagesRefreshed = result.pages_refreshed,
)
@router.get(
"/documents/{document_id}/file-url",
response_model = PreviewFileUrlResponse,
)
def get_document_file_url(
document_id: str,
current_subject: str = Depends(get_current_subject),
) -> PreviewFileUrlResponse:
"""Mint a short-lived signed URL for PDF.js range requests.
The normal bearer-protected `/file` route stays available for blob
fallback; this route keeps bearer access tokens out of query strings.
"""
document_for_subject_or_404(document_id, current_subject)
token, expires_at = _preview_file_token(
subject = current_subject,
document_id = document_id,
)
url = (
f"/api/rag/documents/{quote(document_id, safe = '')}/file-signed"
f"?token={quote(token, safe = '')}"
)
return PreviewFileUrlResponse(url = url, expiresAt = expires_at)
@router.get("/documents/{document_id}/file-signed", response_model = None)
def get_signed_document_file(
document_id: str,
request: Request,
token: str = Query(...),
) -> FileResponse | Response | StreamingResponse:
subject = _subject_from_preview_file_token(document_id, token)
doc_row = document_for_subject_or_404(document_id, subject)
return _serve_document_file_row(
doc_row,
document_id,
request.headers.get("range"),
)
@router.get("/documents/{document_id}/file", response_model = None)
def get_document_file(
document_id: str,
request: Request,
current_subject: str = Depends(get_current_subject),
) -> FileResponse | Response | StreamingResponse:
"""Serve the original uploaded bytes for ``document_id``.
Path resolution per contracts.md §2.1: the route NEVER accepts a
client-supplied filename. ``stored_path`` is DB-issued and we run it
through ``resolve_under_root`` which delegates to ``_assert_contained``
(realpath + symlink/junction-safe). Any escape (symlink, junction,
``..``, absolute outside root) collapses to the second 404.
Content-Type and disposition come from the extension allowlist
(``_preview_file_metadata``). HTML/DOCX/unknown serve as
``attachment`` with safe content-type so an uploaded ``.html`` can
never execute in the app origin (Risk #3 / decisions Q7).
"""
doc_row = document_for_subject_or_404(document_id, current_subject)
return _serve_document_file_row(
doc_row,
document_id,
request.headers.get("range"),
)
@router.post("/search", response_model = SearchResponse)
def search(
payload: SearchRequest,
@ -1092,6 +1636,8 @@ def search(
f"""
SELECT c.id AS chunk_id, c.document_id, c.chunk_index, c.text,
c.page_number, c.kind, c.image_path, c.linked_chunk_id,
c.source_page_index, c.page_char_start, c.page_char_end,
c.line_start, c.line_end,
d.filename
FROM rag_chunks c
JOIN rag_documents d ON d.id = c.document_id
@ -1139,6 +1685,11 @@ def search(
filename = meta.get("filename"),
kind = kind,
image_url = image_url,
source_page_index = meta.get("source_page_index"),
page_char_start = meta.get("page_char_start"),
page_char_end = meta.get("page_char_end"),
line_start = meta.get("line_start"),
line_end = meta.get("line_end"),
)
)
logger.info("RAG search: returning %d hits", len(out))

View file

@ -275,6 +275,12 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
kind TEXT NOT NULL DEFAULT 'text',
image_path TEXT,
linked_chunk_id TEXT,
source_page_index INTEGER,
page_char_start INTEGER,
page_char_end INTEGER,
line_start INTEGER,
line_end INTEGER,
pdf_regions_json TEXT,
UNIQUE(document_id, chunk_index)
)
"""
@ -290,9 +296,38 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.execute("ALTER TABLE rag_chunks ADD COLUMN image_path TEXT")
if "linked_chunk_id" not in chunk_cols:
conn.execute("ALTER TABLE rag_chunks ADD COLUMN linked_chunk_id TEXT")
for column in (
"source_page_index",
"page_char_start",
"page_char_end",
"line_start",
"line_end",
):
if column not in chunk_cols:
conn.execute(f"ALTER TABLE rag_chunks ADD COLUMN {column} INTEGER")
if "pdf_regions_json" not in chunk_cols:
conn.execute("ALTER TABLE rag_chunks ADD COLUMN pdf_regions_json TEXT")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_chunks_document_id ON rag_chunks(document_id)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS rag_document_pages (
document_id TEXT NOT NULL REFERENCES rag_documents(id) ON DELETE CASCADE,
page_index INTEGER NOT NULL,
page_number INTEGER,
text TEXT NOT NULL,
char_count INTEGER NOT NULL DEFAULT 0,
line_count INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
PRIMARY KEY(document_id, page_index)
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_document_pages_document_id "
"ON rag_document_pages(document_id)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS rag_ingestion_jobs (

View file

@ -0,0 +1,269 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for document_for_subject_or_404 and chunk_belongs_to_document.
Authorization rules under test (contracts.md §1 / §2, Risk #1):
- KB documents: KB must exist and KB.owner_user_id must equal current_subject.
- Thread documents: thread must exist in chat_threads; current-Studio single-user
invariant means any authenticated subject can access, BUT the thread row must
exist (a missing thread is 404, not silent grant).
- Missing document or missing KB both collapse to 404.
- KB with NULL owner_user_id is NOT accessible (legacy row guard).
- Both not-found and not-authorized return HTTP 404 with identical detail to
prevent document-existence leaking.
- chunk_belongs_to_document only returns True when chunk.document_id matches.
"""
from __future__ import annotations
import uuid
import pytest
from fastapi import HTTPException
import storage.studio_db as studio_db
from core.rag.authorization import chunk_belongs_to_document, document_for_subject_or_404
# ── Fixtures ──────────────────────────────────────────────────────────
def _reset_db(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
def _uid() -> str:
return str(uuid.uuid4())
def _insert_kb(conn, kb_id: str, owner: str | None = "user-alice") -> None:
conn.execute(
"""
INSERT INTO rag_knowledge_bases (id, name, embedding_model, owner_user_id, created_at)
VALUES (?, ?, ?, ?, ?)
""",
(kb_id, f"KB-{kb_id[:8]}", "bge-small", owner, 1_700_000_000),
)
def _insert_thread(conn, thread_id: str) -> None:
conn.execute(
"""
INSERT INTO chat_threads (id, title, model_type, model_id, archived, created_at)
VALUES (?, ?, ?, ?, ?, ?)
""",
(thread_id, "Test Thread", "base", "llama3", 0, 1_700_000_000),
)
def _insert_kb_doc(conn, doc_id: str, kb_id: str, stored_path: str = "doc.pdf") -> None:
conn.execute(
"""
INSERT INTO rag_documents
(id, kb_id, thread_id, filename, content_type, stored_path, status,
num_chunks, byte_size, created_at)
VALUES (?, ?, NULL, ?, ?, ?, 'completed', 0, 1024, ?)
""",
(doc_id, kb_id, "report.pdf", "application/pdf", stored_path, 1_700_000_000),
)
def _insert_thread_doc(
conn, doc_id: str, thread_id: str, stored_path: str = "doc.txt"
) -> None:
conn.execute(
"""
INSERT INTO rag_documents
(id, kb_id, thread_id, filename, content_type, stored_path, status,
num_chunks, byte_size, created_at)
VALUES (?, NULL, ?, ?, ?, ?, 'completed', 0, 512, ?)
""",
(doc_id, thread_id, "note.txt", "text/plain", stored_path, 1_700_000_000),
)
def _insert_chunk(conn, chunk_id: str, doc_id: str, chunk_index: int = 0) -> None:
conn.execute(
"""
INSERT INTO rag_chunks (id, document_id, chunk_index, text, token_count)
VALUES (?, ?, ?, ?, ?)
""",
(chunk_id, doc_id, chunk_index, "some chunk text", 20),
)
# ── KB-document authorization ─────────────────────────────────────────
def test_kb_doc_correct_owner_returns_row(tmp_path, monkeypatch):
"""KB doc authorized when KB.owner_user_id == current_subject."""
_reset_db(tmp_path, monkeypatch)
doc_id, kb_id = _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_kb_doc(conn, doc_id, kb_id)
row = document_for_subject_or_404(doc_id, "alice")
assert row["id"] == doc_id
def test_kb_doc_wrong_owner_raises_404(tmp_path, monkeypatch):
"""KB doc returns 404 when current_subject != KB.owner_user_id."""
_reset_db(tmp_path, monkeypatch)
doc_id, kb_id = _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_kb_doc(conn, doc_id, kb_id)
with pytest.raises(HTTPException) as exc_info:
document_for_subject_or_404(doc_id, "mallory")
assert exc_info.value.status_code == 404
assert exc_info.value.detail == "Document not found"
def test_kb_doc_null_owner_raises_404(tmp_path, monkeypatch):
"""KB with NULL owner_user_id is not accessible through the helper (legacy guard)."""
_reset_db(tmp_path, monkeypatch)
doc_id, kb_id = _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner=None)
_insert_kb_doc(conn, doc_id, kb_id)
with pytest.raises(HTTPException) as exc_info:
document_for_subject_or_404(doc_id, "alice")
assert exc_info.value.status_code == 404
assert exc_info.value.detail == "Document not found"
def test_kb_doc_missing_kb_raises_404(tmp_path, monkeypatch):
"""Document rows whose KB was deleted collapse to 404.
Insert both KB and doc, then delete the KB (ON DELETE CASCADE removes the doc
too). A subsequent lookup for the doc id must return 404, not 500.
If for some reason the doc row survives (e.g. FK off), the helper must
still 404 because the KB is gone.
"""
_reset_db(tmp_path, monkeypatch)
doc_id, kb_id = _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_kb_doc(conn, doc_id, kb_id)
# Delete the KB — ON DELETE CASCADE should also drop the doc.
conn.execute("DELETE FROM rag_knowledge_bases WHERE id = ?", (kb_id,))
# After cascade deletion the doc_id no longer exists → 404.
with pytest.raises(HTTPException) as exc_info:
document_for_subject_or_404(doc_id, "alice")
assert exc_info.value.status_code == 404
# ── Thread-document authorization (single-user invariant) ─────────────
def test_thread_doc_existing_thread_grants_access(tmp_path, monkeypatch):
"""Thread doc is accessible when thread exists (single-user Studio invariant)."""
_reset_db(tmp_path, monkeypatch)
doc_id, thread_id = _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_thread(conn, thread_id)
_insert_thread_doc(conn, doc_id, thread_id)
row = document_for_subject_or_404(doc_id, "any-authenticated-user")
assert row["id"] == doc_id
def test_thread_doc_nonexistent_thread_raises_404(tmp_path, monkeypatch):
"""A missing thread_id does NOT silently grant access — it must be 404."""
_reset_db(tmp_path, monkeypatch)
doc_id, thread_id = _uid(), _uid()
# Insert doc with a thread_id that has no matching chat_threads row.
with studio_db.get_connection() as conn:
conn.execute(
"""
INSERT INTO rag_documents
(id, kb_id, thread_id, filename, content_type, stored_path, status,
num_chunks, byte_size, created_at)
VALUES (?, NULL, ?, 'x.txt', 'text/plain', 'x.txt', 'completed', 0, 1, ?)
""",
(doc_id, thread_id, 1_700_000_000),
)
with pytest.raises(HTTPException) as exc_info:
document_for_subject_or_404(doc_id, "alice")
assert exc_info.value.status_code == 404
assert exc_info.value.detail == "Document not found"
# ── Missing document ──────────────────────────────────────────────────
def test_missing_document_raises_404(tmp_path, monkeypatch):
"""Completely absent document_id returns 404 with canonical detail."""
_reset_db(tmp_path, monkeypatch)
with pytest.raises(HTTPException) as exc_info:
document_for_subject_or_404("nonexistent-id", "alice")
assert exc_info.value.status_code == 404
assert exc_info.value.detail == "Document not found"
def test_empty_document_id_raises_404(tmp_path, monkeypatch):
"""Empty string document_id raises 404 rather than hitting the DB."""
_reset_db(tmp_path, monkeypatch)
with pytest.raises(HTTPException) as exc_info:
document_for_subject_or_404("", "alice")
assert exc_info.value.status_code == 404
def test_empty_subject_raises_404(tmp_path, monkeypatch):
"""Empty subject raises 404 — cannot authorize without a subject."""
_reset_db(tmp_path, monkeypatch)
doc_id, kb_id = _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_kb_doc(conn, doc_id, kb_id)
with pytest.raises(HTTPException) as exc_info:
document_for_subject_or_404(doc_id, "")
assert exc_info.value.status_code == 404
# ── chunk_belongs_to_document ─────────────────────────────────────────
def test_chunk_belongs_returns_true_for_matching_doc(tmp_path, monkeypatch):
"""chunk_belongs_to_document returns True when chunk.document_id matches."""
_reset_db(tmp_path, monkeypatch)
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_kb_doc(conn, doc_id, kb_id)
_insert_chunk(conn, chunk_id, doc_id)
assert chunk_belongs_to_document(chunk_id, doc_id) is True
def test_chunk_belongs_returns_false_for_wrong_doc(tmp_path, monkeypatch):
"""chunk_belongs_to_document returns False when chunk belongs to a different document."""
_reset_db(tmp_path, monkeypatch)
kb_id = _uid()
doc_a, doc_b, chunk_id = _uid(), _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_kb_doc(conn, doc_a, kb_id, "a.pdf")
_insert_kb_doc(conn, doc_b, kb_id, "b.pdf")
_insert_chunk(conn, chunk_id, doc_a)
# chunk belongs to doc_a — probing with doc_b must return False
assert chunk_belongs_to_document(chunk_id, doc_b) is False
def test_chunk_belongs_returns_false_for_missing_chunk(tmp_path, monkeypatch):
"""chunk_belongs_to_document returns False for a nonexistent chunk_id."""
_reset_db(tmp_path, monkeypatch)
doc_id, kb_id = _uid(), _uid()
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_kb_doc(conn, doc_id, kb_id)
assert chunk_belongs_to_document("ghost-chunk-id", doc_id) is False
def test_chunk_belongs_returns_false_for_empty_inputs(tmp_path, monkeypatch):
"""chunk_belongs_to_document returns False for empty inputs without DB access."""
_reset_db(tmp_path, monkeypatch)
assert chunk_belongs_to_document("", "some-doc") is False
assert chunk_belongs_to_document("some-chunk", "") is False
assert chunk_belongs_to_document("", "") is False

View file

@ -0,0 +1,279 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import queue as queue_module
import uuid
import pytest
import storage.studio_db as studio_db
from core.rag.chunking import chunk_pages, chunk_pages_with_spans
from core.rag.ingestion import (
_JobState,
_insert_chunks_and_collect_for_bm25,
_pump,
_replace_document_pages,
)
from core.rag.parsers import ParsedPage
def _uid() -> str:
return str(uuid.uuid4())
def _token_count(text: str) -> int:
return max(1, len(text.split()))
def test_standard_chunking_records_page_local_char_and_line_spans():
pages = [
ParsedPage(
text="alpha first line\nbeta target line\ngamma final line",
page_number=7,
)
]
chunks = chunk_pages(
pages,
max_tokens=3,
overlap_tokens=0,
token_counter=_token_count,
separators=("\n", " ", ""),
)
target = next(chunk for chunk in chunks if "beta" in chunk.text)
assert target.page_number == 7
assert target.source_page_index == 0
assert target.page_char_start == pages[0].text.index("beta target line")
assert target.page_char_end == target.page_char_start + len("beta target line")
assert target.line_start == 2
assert target.line_end == 2
def test_late_chunking_maps_global_span_back_to_source_page():
pages = [
ParsedPage(text="page one alpha", page_number=1),
ParsedPage(text="page two beta target", page_number=2),
]
_full_doc, chunks, spans = chunk_pages_with_spans(
pages,
max_tokens=4,
overlap_tokens=0,
token_counter=_token_count,
separators=("\n\n", " ", ""),
)
target = next(chunk for chunk in chunks if "beta" in chunk.text)
assert spans[chunks.index(target)][0] >= len(pages[0].text)
assert target.page_number == 2
assert target.source_page_index == 1
assert target.page_char_start is not None
assert target.page_char_end is not None
assert pages[1].text[target.page_char_start : target.page_char_end].strip()
def test_image_chunk_persistence_keeps_page_focus_and_null_text_locators(
tmp_path,
monkeypatch,
):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
from core.rag import ingestion
captured_points: list[dict] = []
monkeypatch.setattr(
ingestion.vector_store,
"upsert_chunks",
lambda _scope, points: captured_points.extend(points),
)
kb_id = _uid()
doc_id = _uid()
with studio_db.get_connection() as conn:
conn.execute(
"""
INSERT INTO rag_knowledge_bases
(id, name, embedding_model, owner_user_id, created_at)
VALUES (?, ?, ?, ?, ?)
""",
(kb_id, "KB", "embedder", "alice", 1_700_000_000),
)
conn.execute(
"""
INSERT INTO rag_documents
(id, kb_id, thread_id, filename, content_type, stored_path, status,
num_chunks, byte_size, created_at)
VALUES (?, ?, NULL, ?, ?, ?, 'completed', 1, 10, ?)
""",
(doc_id, kb_id, "image.pdf", "application/pdf", "image.pdf", 1_700_000_001),
)
_insert_chunks_and_collect_for_bm25(
doc_id,
"kb_scope",
0,
[
{
"text": "",
"token_count": 0,
"page_number": 3,
"kind": "image",
"image_path": str(tmp_path / "img.png"),
}
],
[[0.1, 0.2]],
)
with studio_db.get_connection() as conn:
row = conn.execute(
"""
SELECT page_number, source_page_index, page_char_start,
page_char_end, line_start, line_end
FROM rag_chunks WHERE document_id = ?
""",
(doc_id,),
).fetchone()
assert row["page_number"] == 3
assert row["source_page_index"] is None
assert row["page_char_start"] is None
assert row["page_char_end"] is None
assert row["line_start"] is None
assert row["line_end"] is None
assert captured_points[0]["payload"]["page_number"] == 3
assert captured_points[0]["payload"]["page_char_start"] is None
def test_replace_document_pages_replaces_existing_rows(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
kb_id = _uid()
doc_id = _uid()
with studio_db.get_connection() as conn:
conn.execute(
"""
INSERT INTO rag_knowledge_bases
(id, name, embedding_model, owner_user_id, created_at)
VALUES (?, ?, ?, ?, ?)
""",
(kb_id, "KB", "embedder", "alice", 1_700_000_000),
)
conn.execute(
"""
INSERT INTO rag_documents
(id, kb_id, thread_id, filename, content_type, stored_path, status,
num_chunks, byte_size, created_at)
VALUES (?, ?, NULL, ?, ?, ?, 'completed', 1, 10, ?)
""",
(doc_id, kb_id, "doc.pdf", "application/pdf", "doc.pdf", 1_700_000_001),
)
_replace_document_pages(
doc_id,
[
{
"page_index": 0,
"page_number": 1,
"text": "old page",
"char_count": 8,
"line_count": 1,
}
],
)
_replace_document_pages(
doc_id,
[
{
"page_index": 1,
"page_number": 2,
"text": "new\npage",
"char_count": 8,
"line_count": 2,
}
],
)
with studio_db.get_connection() as conn:
rows = conn.execute(
"""
SELECT page_index, page_number, text, char_count, line_count
FROM rag_document_pages WHERE document_id = ?
""",
(doc_id,),
).fetchall()
assert [dict(row) for row in rows] == [
{
"page_index": 1,
"page_number": 2,
"text": "new\npage",
"char_count": 8,
"line_count": 2,
}
]
class _OneMessageQueue:
def __init__(self, message: dict) -> None:
self.message = message
self.used = False
def get(self, timeout: float) -> dict:
if self.used:
raise queue_module.Empty
self.used = True
return self.message
class _FinishedProcess:
def join(self, timeout: float | None = None) -> None:
return None
def is_alive(self) -> bool:
return False
def terminate(self) -> None:
return None
@pytest.mark.parametrize(
"pages",
[
[
{
"page_index": 0,
"page_number": 1,
"text": "orphan page",
"char_count": 11,
"line_count": 1,
}
],
[],
],
)
def test_document_pages_missing_document_fails_pump_cleanly(
tmp_path,
monkeypatch,
pages,
):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
with studio_db.get_connection():
pass
state = _JobState("job-missing-doc", "missing-doc", "kb_scope")
queue = _OneMessageQueue(
{
"type": "document_pages",
"pages": pages,
}
)
_pump(state, _FinishedProcess(), queue)
assert state.status == "failed"
assert state.error is not None
assert "document was removed before ingestion finished" in state.error

View file

@ -0,0 +1,140 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import uuid
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
import storage.studio_db as studio_db
from auth.authentication import get_current_subject
@pytest.fixture(scope="module")
def app():
import sys
backend_dir = str(Path(__file__).resolve().parent.parent)
if backend_dir not in sys.path:
sys.path.insert(0, backend_dir)
from main import app as _app
return _app
@pytest.fixture
def db_env(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
return tmp_path
def _uid() -> str:
return str(uuid.uuid4())
def _make_client(app, subject: str = "alice"):
app.dependency_overrides[get_current_subject] = lambda: subject
return TestClient(app, raise_server_exceptions=True)
def _clear_overrides(app):
app.dependency_overrides.clear()
def _insert_kb(conn, kb_id: str, owner: str = "alice") -> None:
conn.execute(
"INSERT INTO rag_knowledge_bases "
"(id, name, embedding_model, owner_user_id, created_at) "
"VALUES (?, ?, ?, ?, ?)",
(kb_id, f"KB-{kb_id[:6]}", "bge-small", owner, 1_700_000_000),
)
def _insert_doc(conn, doc_id: str, kb_id: str, stored_path: str, filename: str) -> None:
conn.execute(
"INSERT INTO rag_documents "
"(id, kb_id, thread_id, filename, content_type, stored_path, status, "
"num_chunks, byte_size, created_at) "
"VALUES (?, ?, NULL, ?, 'text/plain', ?, 'completed', 1, 64, ?)",
(doc_id, kb_id, filename, stored_path, 1_700_000_000),
)
def _insert_chunk(conn, chunk_id: str, doc_id: str, text: str) -> None:
conn.execute(
"INSERT INTO rag_chunks "
"(id, document_id, chunk_index, text, token_count, page_number) "
"VALUES (?, ?, 0, ?, 5, NULL)",
(chunk_id, doc_id, text),
)
def test_backfill_preserves_ids_and_updates_unique_locator(app, db_env, monkeypatch):
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
stored = db_env / "rag" / "uploads" / "paper.txt"
stored.parent.mkdir(parents=True, exist_ok=True)
stored.write_text("Intro line\nUnique quote here.\nEnd.", encoding="utf-8")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored), "paper.txt")
_insert_chunk(conn, chunk_id, doc_id, "Unique quote here.")
client = _make_client(app, "alice")
try:
resp = client.post(f"/api/rag/documents/{doc_id}/locators/backfill")
target_resp = client.get(
f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}"
)
finally:
_clear_overrides(app)
assert resp.status_code == 200
body = resp.json()
assert body["documentId"] == doc_id
assert body["matched"] == 1
assert body["ambiguous"] == 0
target = target_resp.json()
assert target["documentId"] == doc_id
assert target["chunkId"] == chunk_id
assert target["sourcePageIndex"] == 0
assert target["lineStart"] == 2
assert target["pageCharStart"] in (len("Intro line\n"), len("Intro line\r\n"))
def test_backfill_leaves_ambiguous_matches_null(app, db_env, monkeypatch):
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
stored = db_env / "rag" / "uploads" / "paper.txt"
stored.parent.mkdir(parents=True, exist_ok=True)
stored.write_text("Repeat me.\nOther text.\nRepeat me.", encoding="utf-8")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored), "paper.txt")
_insert_chunk(conn, chunk_id, doc_id, "Repeat me.")
client = _make_client(app, "alice")
try:
resp = client.post(f"/api/rag/documents/{doc_id}/locators/backfill")
target_resp = client.get(
f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}"
)
finally:
_clear_overrides(app)
assert resp.status_code == 200
body = resp.json()
assert body["matched"] == 0
assert body["ambiguous"] == 1
target = target_resp.json()
assert target["sourcePageIndex"] is None
assert target["pageCharStart"] is None
assert target["lineStart"] is None

View file

@ -0,0 +1,87 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import uuid
import storage.studio_db as studio_db
def _uid() -> str:
return str(uuid.uuid4())
def test_locator_schema_is_additive_and_nullable(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
with studio_db.get_connection() as conn:
chunk_cols = {
row["name"] for row in conn.execute("PRAGMA table_info(rag_chunks)")
}
assert {
"source_page_index",
"page_char_start",
"page_char_end",
"line_start",
"line_end",
}.issubset(chunk_cols)
page_cols = {
row["name"]
for row in conn.execute("PRAGMA table_info(rag_document_pages)")
}
assert {
"document_id",
"page_index",
"page_number",
"text",
"char_count",
"line_count",
}.issubset(page_cols)
kb_id = _uid()
doc_id = _uid()
chunk_id = _uid()
conn.execute(
"""
INSERT INTO rag_knowledge_bases
(id, name, embedding_model, owner_user_id, created_at)
VALUES (?, ?, ?, ?, ?)
""",
(kb_id, "KB", "embedder", "alice", 1_700_000_000),
)
conn.execute(
"""
INSERT INTO rag_documents
(id, kb_id, thread_id, filename, content_type, stored_path, status,
num_chunks, byte_size, created_at)
VALUES (?, ?, NULL, ?, ?, ?, 'completed', 1, 10, ?)
""",
(doc_id, kb_id, "old.pdf", "application/pdf", "old.pdf", 1_700_000_001),
)
conn.execute(
"""
INSERT INTO rag_chunks
(id, document_id, chunk_index, text, token_count, page_number)
VALUES (?, ?, 0, ?, 3, 1)
""",
(chunk_id, doc_id, "legacy chunk"),
)
row = conn.execute(
"""
SELECT source_page_index, page_char_start, page_char_end,
line_start, line_end
FROM rag_chunks WHERE id = ?
""",
(chunk_id,),
).fetchone()
assert dict(row) == {
"source_page_index": None,
"page_char_start": None,
"page_char_end": None,
"line_start": None,
"line_end": None,
}

View file

@ -0,0 +1,593 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for GET /api/rag/documents/{id}/preview-target and /file.
Acceptance criteria covered (contracts.md §1, §2, PLAN.md T1/T2, Risk #1-3):
/preview-target:
- 200 with chunk data when chunk_id present and belongs to doc.
- 200 with all-null chunk fields when chunk_id absent (document-row preview).
- 404 when document missing (collapsed existence + auth).
- 404 when chunk_id does not belong to document_id (cross-doc probe collapsed).
- 401 when no bearer token.
/file:
- 200 with correct Content-Type and nosniff header.
- X-Content-Type-Options: nosniff present on every 200.
- Cache-Control: private present on every 200.
- HTML extension served as text/plain + attachment (Risk #3).
- DOCX extension served with attachment disposition.
- 404 when document missing or wrong subject.
- 404 when file deleted from disk (DB row exists, subject authorized).
- Outside-root stored_path returns 404 (path containment, Risk #2).
Auth is injected via dependency override (mock at the boundary, not the
implementation target). We do NOT mock document_for_subject_or_404 itself.
"""
from __future__ import annotations
import os
import uuid
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
import storage.studio_db as studio_db
from auth.authentication import get_current_subject
# ── App import (deferred to avoid import-time side-effects) ───────────
@pytest.fixture(scope="module")
def app():
import sys
backend_dir = str(Path(__file__).resolve().parent.parent)
if backend_dir not in sys.path:
sys.path.insert(0, backend_dir)
from main import app as _app
return _app
# ── Test-level DB + auth fixtures ─────────────────────────────────────
@pytest.fixture
def db_env(tmp_path, monkeypatch):
"""Point studio_db at a fresh temp DB for each test."""
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
return tmp_path
def _uid() -> str:
return str(uuid.uuid4())
def _make_client(app, subject: str = "alice"):
"""Return a TestClient with get_current_subject overridden to return subject."""
app.dependency_overrides[get_current_subject] = lambda: subject
client = TestClient(app, raise_server_exceptions=True)
return client
def _clear_overrides(app):
app.dependency_overrides.clear()
def _insert_kb(conn, kb_id: str, owner: str = "alice") -> None:
conn.execute(
"INSERT INTO rag_knowledge_bases (id, name, embedding_model, owner_user_id, created_at) "
"VALUES (?, ?, ?, ?, ?)",
(kb_id, f"KB-{kb_id[:6]}", "bge-small", owner, 1_700_000_000),
)
def _insert_doc(
conn,
doc_id: str,
kb_id: str,
stored_path: str,
filename: str = "report.pdf",
content_type: str | None = "application/pdf",
status: str = "completed",
) -> None:
conn.execute(
"INSERT INTO rag_documents "
"(id, kb_id, thread_id, filename, content_type, stored_path, status, "
"num_chunks, byte_size, created_at) "
"VALUES (?, ?, NULL, ?, ?, ?, ?, 0, 1024, ?)",
(doc_id, kb_id, filename, content_type, stored_path, status, 1_700_000_000),
)
def _insert_chunk(
conn,
chunk_id: str,
doc_id: str,
text: str = "The margin rose to 18.2% in Q3.",
page_number: int | None = 7,
chunk_index: int = 14,
) -> None:
conn.execute(
"INSERT INTO rag_chunks "
"(id, document_id, chunk_index, text, token_count, page_number) "
"VALUES (?, ?, ?, ?, ?, ?)",
(chunk_id, doc_id, chunk_index, text, 30, page_number),
)
# ── /preview-target tests ─────────────────────────────────────────────
class TestPreviewTarget:
def test_with_chunk_id_returns_full_metadata(self, app, db_env, monkeypatch):
"""GET /preview-target?chunk_id=<id> returns page + snippet when chunk valid."""
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
stored = db_env / "rag" / "uploads" / "report.pdf"
stored.parent.mkdir(parents=True, exist_ok=True)
stored.write_bytes(b"%PDF-1.4 dummy")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored))
_insert_chunk(conn, chunk_id, doc_id, page_number=7, chunk_index=14)
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}")
finally:
_clear_overrides(app)
assert resp.status_code == 200
body = resp.json()
assert body["documentId"] == doc_id
assert body["chunkId"] == chunk_id
assert body["targetPage"] == 7
assert body["chunkIndex"] == 14
assert body["snippet"] is not None and len(body["snippet"]) > 0
assert body["mediaKind"] == "pdf"
def test_preview_target_returns_pdf_regions_when_present(self, app, db_env, monkeypatch):
"""Chunk preview includes only stored confident PDF regions."""
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
stored = db_env / "rag" / "uploads" / "report.pdf"
stored.parent.mkdir(parents=True, exist_ok=True)
stored.write_bytes(b"%PDF-1.4 dummy")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored))
_insert_chunk(conn, chunk_id, doc_id, page_number=7, chunk_index=14)
conn.execute(
"""
UPDATE rag_chunks
SET pdf_regions_json = ?
WHERE id = ?
""",
(
'[{"pageIndex":6,"pageNumber":7,"x":0.1,"y":0.2,'
'"width":0.3,"height":0.04,"confidence":"exact",'
'"source":"pymupdf-search"}]',
chunk_id,
),
)
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}")
finally:
_clear_overrides(app)
assert resp.status_code == 200
body = resp.json()
assert body["pdfRegions"] == [
{
"pageIndex": 6,
"pageNumber": 7,
"x": 0.1,
"y": 0.2,
"width": 0.3,
"height": 0.04,
"confidence": "exact",
"source": "pymupdf-search",
}
]
def test_without_chunk_id_returns_all_null_chunk_fields(self, app, db_env, monkeypatch):
"""GET /preview-target without chunk_id returns metadata-only (decision Q2)."""
doc_id, kb_id = _uid(), _uid()
stored = db_env / "rag" / "uploads" / "annual.pdf"
stored.parent.mkdir(parents=True, exist_ok=True)
stored.write_bytes(b"%PDF-1.4 dummy")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored))
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/preview-target")
finally:
_clear_overrides(app)
assert resp.status_code == 200
body = resp.json()
# All chunk fields MUST be null — UI must not guess a first chunk.
assert body["chunkId"] is None
assert body["chunkIndex"] is None
assert body["targetPage"] is None
assert body["snippet"] is None
assert body["kind"] is None
assert body["imageUrl"] is None
assert body["documentId"] == doc_id
def test_missing_document_returns_404(self, app, db_env, monkeypatch):
"""Nonexistent document_id returns 404 to both existence and auth probes."""
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{_uid()}/preview-target")
finally:
_clear_overrides(app)
assert resp.status_code == 404
assert resp.json()["detail"] == "Document not found"
def test_wrong_subject_returns_404(self, app, db_env, monkeypatch):
"""Document owned by alice returns 404 when accessed by mallory."""
doc_id, kb_id = _uid(), _uid()
stored = db_env / "rag" / "uploads" / "secret.pdf"
stored.parent.mkdir(parents=True, exist_ok=True)
stored.write_bytes(b"data")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_doc(conn, doc_id, kb_id, str(stored))
client = _make_client(app, "mallory")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/preview-target")
finally:
_clear_overrides(app)
assert resp.status_code == 404
def test_cross_doc_chunk_id_returns_404(self, app, db_env, monkeypatch):
"""chunk_id from a different document returns 404 — not 400 (opaque)."""
kb_id = _uid()
doc_a, doc_b = _uid(), _uid()
chunk_a = _uid()
stored_a = db_env / "rag" / "uploads" / "a.pdf"
stored_b = db_env / "rag" / "uploads" / "b.pdf"
stored_a.parent.mkdir(parents=True, exist_ok=True)
stored_a.write_bytes(b"data")
stored_b.write_bytes(b"data")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_a, kb_id, str(stored_a), "a.pdf")
_insert_doc(conn, doc_b, kb_id, str(stored_b), "b.pdf")
_insert_chunk(conn, chunk_a, doc_a)
client = _make_client(app, "alice")
try:
# Probe doc_b with chunk_a (which belongs to doc_a)
resp = client.get(
f"/api/rag/documents/{doc_b}/preview-target?chunk_id={chunk_a}"
)
finally:
_clear_overrides(app)
# Must be 404, NOT 200 with doc_a's chunk data
assert resp.status_code == 404
def test_unauthenticated_returns_401(self, app, db_env, monkeypatch):
"""No bearer token → 401."""
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
# No override — let the real dependency raise
client = TestClient(app, raise_server_exceptions=False)
resp = client.get(f"/api/rag/documents/{_uid()}/preview-target")
assert resp.status_code == 401
# ── /file tests ───────────────────────────────────────────────────────
class TestFileRoute:
def test_pdf_200_with_correct_headers(self, app, db_env, monkeypatch):
"""GET /file for a PDF returns 200 with nosniff, Cache-Control, inline disposition."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
stored = uploads / "annual.pdf"
stored.write_bytes(b"%PDF-1.4\n%%EOF")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored), "annual.pdf", "application/pdf")
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/file")
finally:
_clear_overrides(app)
assert resp.status_code == 200
assert resp.headers.get("x-content-type-options") == "nosniff"
assert "private" in (resp.headers.get("cache-control") or "")
ct = resp.headers.get("content-type", "")
assert "pdf" in ct.lower()
def test_signed_file_url_supports_range_without_bearer_query(self, app, db_env, monkeypatch):
"""Short-lived signed URL is redeemable without Authorization and supports ranges."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
stored = uploads / "annual.pdf"
stored.write_bytes(b"%PDF-1.4\n%%EOF")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
monkeypatch.setattr("routes.rag.get_jwt_secret", lambda subject: "test-secret")
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored), "annual.pdf", "application/pdf")
client = _make_client(app, "alice")
try:
url_resp = client.get(f"/api/rag/documents/{doc_id}/file-url")
assert url_resp.status_code == 200
signed_url = url_resp.json()["url"]
assert "Bearer" not in signed_url
assert "Authorization" not in signed_url
file_resp = client.get(signed_url, headers={"Range": "bytes=0-3"})
finally:
_clear_overrides(app)
assert file_resp.status_code == 206
assert file_resp.content == b"%PDF"
assert file_resp.headers.get("content-range") == f"bytes 0-3/{stored.stat().st_size}"
assert file_resp.headers.get("accept-ranges") == "bytes"
assert file_resp.headers.get("x-content-type-options") == "nosniff"
def test_signed_file_route_rejects_forged_token(self, app, db_env, monkeypatch):
"""Signed file route is not public without a valid preview token."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
stored = uploads / "annual.pdf"
stored.write_bytes(b"%PDF-1.4\n%%EOF")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
monkeypatch.setattr("routes.rag.get_jwt_secret", lambda subject: "test-secret")
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored), "annual.pdf", "application/pdf")
client = TestClient(app, raise_server_exceptions=False)
resp = client.get(f"/api/rag/documents/{doc_id}/file-signed?token=bogus")
assert resp.status_code == 401
def test_html_file_served_as_text_plain_with_attachment(self, app, db_env, monkeypatch):
"""HTML uploads must be served as text/plain + attachment (Risk #3 — no XSS)."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
stored = uploads / "malicious.html"
stored.write_bytes(b"<script>alert(1)</script>")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored), "malicious.html", "text/html")
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/file")
finally:
_clear_overrides(app)
assert resp.status_code == 200
ct = resp.headers.get("content-type", "").lower()
# MUST NOT be text/html — must be text/plain
assert "text/html" not in ct, f"HTML executed inline! content-type={ct}"
assert "text/plain" in ct
disp = resp.headers.get("content-disposition", "").lower()
assert "attachment" in disp, f"HTML not forced to attachment: {disp}"
assert resp.headers.get("x-content-type-options") == "nosniff"
def test_docx_served_as_attachment(self, app, db_env, monkeypatch):
"""DOCX files must be served with Content-Disposition: attachment."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
stored = uploads / "report.docx"
stored.write_bytes(b"PK\x03\x04fake-docx")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored), "report.docx",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document")
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/file")
finally:
_clear_overrides(app)
assert resp.status_code == 200
disp = resp.headers.get("content-disposition", "").lower()
assert "attachment" in disp
def test_missing_document_returns_404(self, app, db_env, monkeypatch):
"""Nonexistent document returns 404."""
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{_uid()}/file")
finally:
_clear_overrides(app)
assert resp.status_code == 404
def test_wrong_subject_returns_404(self, app, db_env, monkeypatch):
"""Document accessible to alice is 404 for mallory (auth-collapse)."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
stored = uploads / "private.pdf"
stored.write_bytes(b"data")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_doc(conn, doc_id, kb_id, str(stored))
client = _make_client(app, "mallory")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/file")
finally:
_clear_overrides(app)
assert resp.status_code == 404
def test_deleted_file_returns_404_with_doc_file_not_found(self, app, db_env, monkeypatch):
"""File gone from disk returns 404 with 'Document file not found' detail."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
stored = uploads / "gone.pdf"
stored.write_bytes(b"data")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored))
# Delete the file after inserting the row
stored.unlink()
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/file")
finally:
_clear_overrides(app)
assert resp.status_code == 404
detail = resp.json().get("detail", "")
assert "file not found" in detail.lower() or "not found" in detail.lower()
def test_outside_root_stored_path_returns_404(self, app, db_env, monkeypatch, tmp_path):
"""stored_path outside rag_uploads_root returns 404 — path containment (Risk #2)."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
# A legitimate-looking path that is outside the RAG uploads root
outside = tmp_path / "etc" / "passwd"
outside.parent.mkdir(parents=True, exist_ok=True)
outside.write_bytes(b"root:x:0:0")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
# Insert with stored_path pointing outside root
conn.execute(
"INSERT INTO rag_documents "
"(id, kb_id, thread_id, filename, content_type, stored_path, status, "
"num_chunks, byte_size, created_at) "
"VALUES (?, ?, NULL, 'passwd', 'text/plain', ?, 'completed', 0, 10, ?)",
(doc_id, kb_id, str(outside), 1_700_000_000),
)
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/file")
finally:
_clear_overrides(app)
# Must NOT serve the file — containment violation must return 404
assert resp.status_code == 404
def test_nosniff_and_cache_headers_on_txt_file(self, app, db_env, monkeypatch):
"""Safety headers present on every 200 response, including plain text."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
uploads.mkdir(parents=True, exist_ok=True)
stored = uploads / "notes.txt"
stored.write_bytes(b"hello world")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id)
_insert_doc(conn, doc_id, kb_id, str(stored), "notes.txt", "text/plain")
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/documents/{doc_id}/file")
finally:
_clear_overrides(app)
assert resp.status_code == 200
assert resp.headers.get("x-content-type-options") == "nosniff"
cc = resp.headers.get("cache-control", "")
assert "private" in cc
# ── /images tests ─────────────────────────────────────────────────────
class TestImageRoute:
def test_image_route_wrong_subject_returns_404(self, app, db_env, monkeypatch):
"""Extracted images require the same document authorization as /file."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
images = uploads / "images" / doc_id
images.mkdir(parents=True, exist_ok=True)
image = images / "figure.png"
image.write_bytes(b"\x89PNG\r\n\x1a\n")
stored = uploads / "report.pdf"
stored.write_bytes(b"%PDF-1.4")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_doc(conn, doc_id, kb_id, str(stored), "report.pdf")
client = _make_client(app, "mallory")
try:
resp = client.get(f"/api/rag/images/{doc_id}/figure.png")
finally:
_clear_overrides(app)
assert resp.status_code == 404
def test_image_route_authorized_subject_gets_image(self, app, db_env, monkeypatch):
"""Authorized subject can still fetch an extracted image."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
images = uploads / "images" / doc_id
images.mkdir(parents=True, exist_ok=True)
image = images / "figure.png"
image.write_bytes(b"\x89PNG\r\n\x1a\n")
stored = uploads / "report.pdf"
stored.write_bytes(b"%PDF-1.4")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner="alice")
_insert_doc(conn, doc_id, kb_id, str(stored), "report.pdf")
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/images/{doc_id}/figure.png")
finally:
_clear_overrides(app)
assert resp.status_code == 200
assert resp.content.startswith(b"\x89PNG")

View file

@ -0,0 +1,136 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import uuid
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
import storage.studio_db as studio_db
from auth.authentication import get_current_subject
@pytest.fixture(scope="module")
def app():
import sys
backend_dir = str(Path(__file__).resolve().parent.parent)
if backend_dir not in sys.path:
sys.path.insert(0, backend_dir)
from main import app as _app
return _app
@pytest.fixture
def db_env(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
return tmp_path
def _uid() -> str:
return str(uuid.uuid4())
def _make_client(app, subject: str = "alice"):
app.dependency_overrides[get_current_subject] = lambda: subject
return TestClient(app, raise_server_exceptions=True)
def _clear_overrides(app):
app.dependency_overrides.clear()
def _seed_doc(conn, doc_id: str, kb_id: str, stored_path: str) -> None:
conn.execute(
"""
INSERT INTO rag_knowledge_bases
(id, name, embedding_model, owner_user_id, created_at)
VALUES (?, ?, ?, ?, ?)
""",
(kb_id, "KB", "embedder", "alice", 1_700_000_000),
)
conn.execute(
"""
INSERT INTO rag_documents
(id, kb_id, thread_id, filename, content_type, stored_path, status,
num_chunks, byte_size, created_at)
VALUES (?, ?, NULL, ?, ?, ?, 'completed', 1, 10, ?)
""",
(doc_id, kb_id, "report.pdf", "application/pdf", stored_path, 1_700_000_001),
)
def test_preview_target_returns_nullable_locator_fields(app, db_env):
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
stored = db_env / "rag" / "uploads" / "report.pdf"
stored.parent.mkdir(parents=True, exist_ok=True)
stored.write_bytes(b"%PDF-1.4")
with studio_db.get_connection() as conn:
_seed_doc(conn, doc_id, kb_id, str(stored))
conn.execute(
"""
INSERT INTO rag_chunks
(id, document_id, chunk_index, text, token_count, page_number,
source_page_index, page_char_start, page_char_end, line_start,
line_end)
VALUES (?, ?, 2, ?, 8, 4, 3, 20, 52, 6, 7)
""",
(chunk_id, doc_id, "highlight me"),
)
client = _make_client(app)
try:
resp = client.get(
f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}"
)
finally:
_clear_overrides(app)
assert resp.status_code == 200
body = resp.json()
assert body["sourcePageIndex"] == 3
assert body["pageCharStart"] == 20
assert body["pageCharEnd"] == 52
assert body["lineStart"] == 6
assert body["lineEnd"] == 7
def test_preview_target_old_null_locator_rows_still_work(app, db_env):
doc_id, kb_id, chunk_id = _uid(), _uid(), _uid()
stored = db_env / "rag" / "uploads" / "legacy.pdf"
stored.parent.mkdir(parents=True, exist_ok=True)
stored.write_bytes(b"%PDF-1.4")
with studio_db.get_connection() as conn:
_seed_doc(conn, doc_id, kb_id, str(stored))
conn.execute(
"""
INSERT INTO rag_chunks
(id, document_id, chunk_index, text, token_count, page_number)
VALUES (?, ?, 0, ?, 4, 1)
""",
(chunk_id, doc_id, "legacy"),
)
client = _make_client(app)
try:
resp = client.get(
f"/api/rag/documents/{doc_id}/preview-target?chunk_id={chunk_id}"
)
finally:
_clear_overrides(app)
assert resp.status_code == 200
body = resp.json()
assert body["snippet"] == "legacy"
assert body["sourcePageIndex"] is None
assert body["pageCharStart"] is None
assert body["pageCharEnd"] is None
assert body["lineStart"] is None
assert body["lineEnd"] is None

View file

@ -0,0 +1,243 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for durable source identity in tool XML output (PLAN.md T3, contracts §3).
Acceptance criteria:
- _format_hits_for_llm emits document_id and chunk_id attributes on <chunk> elements.
- The visible citation id (id="N") is a per-call counter, NOT the backend chunk UUID.
- Same-filename documents in different KB slots remain distinguishable by document_id.
- Hits without a matching DB row (lookup miss) are silently dropped not emitted
with wrong IDs.
- Legacy hits (no document_id in hit dict) still render without crashing.
"""
from __future__ import annotations
import re
import uuid
from xml.etree import ElementTree
import pytest
from core.rag.tool import _format_hits_for_llm
# ── Helpers ───────────────────────────────────────────────────────────
def _uid() -> str:
return str(uuid.uuid4())
def _hit(
*,
chunk_id: str,
document_id: str,
filename: str = "report.pdf",
text: str = "some text",
page_number: int | None = 3,
chunk_index: int = 0,
score: float = 0.85,
) -> dict:
"""Build a flat hit dict as _format_hits_for_llm expects."""
return {
"chunk_id": chunk_id,
"document_id": document_id,
"filename": filename,
"text": text,
"page_number": page_number,
"chunk_index": chunk_index,
"score": score,
"dense_score": score,
"token_count": 20,
"kind": "text",
"image_path": None,
}
def _parse_chunks(xml_output: str) -> list[dict]:
"""Parse <chunk ...> elements from the multi-block tool output."""
chunks = []
# Each block is wrapped in <chunk ...>\n...\n</chunk>; parse directly.
for match in re.finditer(r"<chunk\s([^>]*)>", xml_output):
attrs_raw = match.group(1)
# Quick attribute parser for "key="value"" pairs.
attrs: dict = {}
for m in re.finditer(r'(\w+)="([^"]*)"', attrs_raw):
attrs[m.group(1)] = m.group(2)
chunks.append(attrs)
return chunks
# ── Tests: durable IDs present in XML ────────────────────────────────
def test_format_hits_emits_document_id_and_chunk_id():
"""T3: tool XML <chunk> must carry document_id and chunk_id attributes."""
chunk_id, doc_id = _uid(), _uid()
hits = [_hit(chunk_id=chunk_id, document_id=doc_id)]
output = _format_hits_for_llm(hits)
chunks = _parse_chunks(output)
assert len(chunks) == 1, output
assert chunks[0]["document_id"] == doc_id
assert chunks[0]["chunk_id"] == chunk_id
def test_citation_id_is_sequential_counter_not_uuid():
"""Visible id='N' is a 1-based counter — never equal to the backend chunk UUID."""
chunk_id, doc_id = _uid(), _uid()
hits = [_hit(chunk_id=chunk_id, document_id=doc_id)]
output = _format_hits_for_llm(hits, start_id=0)
chunks = _parse_chunks(output)
visible_id = chunks[0]["id"]
# Must be a small integer string, NOT the UUID
assert visible_id == "1", f"expected '1' got {visible_id!r}"
assert visible_id != chunk_id
def test_citation_ids_are_globally_sequential_across_calls():
"""start_id offset ensures IDs stay unique across multiple tool calls per turn."""
hits_call1 = [_hit(chunk_id=_uid(), document_id=_uid(), filename="a.pdf")]
hits_call2 = [
_hit(chunk_id=_uid(), document_id=_uid(), filename="b.pdf"),
_hit(chunk_id=_uid(), document_id=_uid(), filename="c.pdf"),
]
out1 = _format_hits_for_llm(hits_call1, start_id=0)
out2 = _format_hits_for_llm(hits_call2, start_id=1)
chunks1 = _parse_chunks(out1)
chunks2 = _parse_chunks(out2)
assert chunks1[0]["id"] == "1"
assert chunks2[0]["id"] == "2"
assert chunks2[1]["id"] == "3"
# No id overlap
all_ids = {c["id"] for c in chunks1 + chunks2}
assert len(all_ids) == 3
def test_same_filename_docs_have_distinct_document_ids():
"""Two docs with the same filename route to distinct document_id values (Risk #4)."""
filename = "annual-report.pdf"
chunk_a, doc_a = _uid(), _uid()
chunk_b, doc_b = _uid(), _uid()
hits = [
_hit(chunk_id=chunk_a, document_id=doc_a, filename=filename),
_hit(chunk_id=chunk_b, document_id=doc_b, filename=filename),
]
output = _format_hits_for_llm(hits)
chunks = _parse_chunks(output)
assert len(chunks) == 2
# Both use the same filename but MUST have distinct document_id values
assert chunks[0]["document_id"] != chunks[1]["document_id"]
assert chunks[0]["document_id"] == doc_a
assert chunks[1]["document_id"] == doc_b
def test_same_filename_docs_have_distinct_citation_ids():
"""Same-filename docs in the same turn still get distinct visible [N] ids."""
filename = "notes.pdf"
chunk_a, doc_a = _uid(), _uid()
chunk_b, doc_b = _uid(), _uid()
hits = [
_hit(chunk_id=chunk_a, document_id=doc_a, filename=filename),
_hit(chunk_id=chunk_b, document_id=doc_b, filename=filename),
]
output = _format_hits_for_llm(hits)
chunks = _parse_chunks(output)
citation_ids = {c["id"] for c in chunks}
assert len(citation_ids) == 2, f"citation IDs not unique: {chunks}"
def test_empty_hits_returns_no_chunks_message():
"""Empty hit list returns the 'no matching chunks' message, not broken XML."""
output = _format_hits_for_llm([])
chunks = _parse_chunks(output)
assert len(chunks) == 0
assert "no matching chunks" in output.lower() or "no matching" in output.lower()
def test_page_number_attribute_present_when_page_exists():
"""page attribute is emitted when page_number is not None."""
chunk_id, doc_id = _uid(), _uid()
hits = [_hit(chunk_id=chunk_id, document_id=doc_id, page_number=5)]
output = _format_hits_for_llm(hits)
chunks = _parse_chunks(output)
assert chunks[0].get("page") == "5"
def test_page_number_attribute_absent_when_null():
"""page attribute is omitted when page_number is None."""
chunk_id, doc_id = _uid(), _uid()
hits = [_hit(chunk_id=chunk_id, document_id=doc_id, page_number=None)]
output = _format_hits_for_llm(hits)
chunks = _parse_chunks(output)
assert "page" not in chunks[0], f"unexpected page attr: {chunks[0]}"
def test_locator_attributes_are_additive_when_present():
"""T10: tool XML carries nullable locator metadata without changing visible ids."""
chunk_id, doc_id = _uid(), _uid()
hit = _hit(chunk_id=chunk_id, document_id=doc_id, page_number=5)
hit.update(
{
"source_page_index": 4,
"page_char_start": 11,
"page_char_end": 42,
"line_start": 2,
"line_end": 3,
}
)
output = _format_hits_for_llm([hit])
chunk = _parse_chunks(output)[0]
assert chunk["id"] == "1"
assert chunk["chunk_id"] == chunk_id
assert chunk["source_page_index"] == "4"
assert chunk["page_char_start"] == "11"
assert chunk["page_char_end"] == "42"
assert chunk["line_start"] == "2"
assert chunk["line_end"] == "3"
def test_xml_special_chars_in_filename_escaped():
"""Filename with XML special chars does not break the chunk element."""
chunk_id, doc_id = _uid(), _uid()
hits = [
_hit(
chunk_id=chunk_id,
document_id=doc_id,
filename='report <2025> "final" & draft.pdf',
)
]
output = _format_hits_for_llm(hits)
# The output must parse cleanly (no unescaped < or " in attrs)
chunks = _parse_chunks(output)
assert len(chunks) == 1
# source attribute should have the filename escaped
source_attr = chunks[0].get("source", "")
assert "<" not in source_attr and '"' not in source_attr
def test_multiple_hits_carry_independent_ids():
"""Three hits each carry their own distinct chunk_id and document_id."""
hit_data = [
(_uid(), _uid()),
(_uid(), _uid()),
(_uid(), _uid()),
]
hits = [
_hit(chunk_id=cid, document_id=did, filename=f"doc{i}.pdf")
for i, (cid, did) in enumerate(hit_data)
]
output = _format_hits_for_llm(hits)
chunks = _parse_chunks(output)
assert len(chunks) == 3
emitted_chunk_ids = {c["chunk_id"] for c in chunks}
emitted_doc_ids = {c["document_id"] for c in chunks}
expected_chunk_ids = {cid for cid, _ in hit_data}
expected_doc_ids = {did for _, did in hit_data}
assert emitted_chunk_ids == expected_chunk_ids
assert emitted_doc_ids == expected_doc_ids

View file

@ -38,7 +38,7 @@
},
"overrides": [
{
"include": ["vite.config.ts", "eslint.config.js"],
"include": ["vite.config.ts", "vitest.config.ts", "eslint.config.js"],
"linter": {
"rules": {
"correctness": { "noNodejsModules": "off" },

File diff suppressed because it is too large Load diff

View file

@ -12,6 +12,8 @@
"lint": "eslint .",
"preview": "vite preview",
"typecheck": "tsc -b --pretty false",
"test": "vitest run",
"test:watch": "vitest",
"biome:check": "biome check .",
"biome:fix": "biome check . --write"
},
@ -46,6 +48,7 @@
"@tauri-apps/plugin-process": "^2.3.1",
"@tauri-apps/plugin-updater": "^2.10.1",
"@toolwind/corner-shape": "^0.0.8-3",
"@types/event-source-polyfill": "1.0.5",
"@xyflow/react": "^12.10.0",
"assistant-stream": "0.3.12",
"canvas-confetti": "^1.9.4",
@ -53,6 +56,7 @@
"clsx": "^2.1.1",
"cmdk": "^1.1.1",
"dexie": "^4.3.0",
"event-source-polyfill": "1.0.31",
"fflate": "0.8.3",
"js-yaml": "^4.1.1",
"katex": "^0.16.28",
@ -65,6 +69,7 @@
"react": "^19.2.4",
"react-day-picker": "^9.13.2",
"react-dom": "^19.2.4",
"react-pdf": "^10.4.1",
"react-resizable-panels": "^4.6.4",
"recharts": "3.7.0",
"shadcn": "^4.2.0",
@ -85,10 +90,14 @@
"devDependencies": {
"@biomejs/biome": "^1.9.4",
"@eslint/js": "^9.39.1",
"@testing-library/dom": "^10.4.1",
"@testing-library/jest-dom": "^6.9.1",
"@testing-library/react": "^16.3.2",
"@testing-library/user-event": "^14.6.1",
"@types/canvas-confetti": "^1.9.0",
"@types/js-yaml": "^4.0.9",
"@types/node-forge": "^1.3.14",
"@types/node": "^25.5.2",
"@types/node-forge": "^1.3.14",
"@types/react": "^19.2.5",
"@types/react-dom": "^19.2.3",
"@vitejs/plugin-react": "^6.0.1",
@ -96,8 +105,10 @@
"eslint-plugin-react-hooks": "^7.0.1",
"eslint-plugin-react-refresh": "^0.5.2",
"globals": "^17.4.0",
"jsdom": "^29.1.1",
"typescript": "~5.9.3",
"typescript-eslint": "^8.55.0",
"vite": "^8.0.1"
"vite": "^8.0.1",
"vitest": "^4.1.7"
}
}

View file

@ -0,0 +1,125 @@
/**
* Tests for chat-adapter XML parsing contracts §3 / §4, T3.
*
* Coverage:
* - New XML with document_id + chunk_id attributes citationId, documentId, backendChunkId populated.
* - Legacy XML without durable IDs citationId populated, documentId/backendChunkId absent.
* - Same visible [N] across turns does NOT imply same backendChunkId.
* - Same filename in two chunks distinct documentId values preserved.
* - Missing attributes degrade gracefully no throw.
* - citationId is always the visible "N" counter, never the UUID.
*/
import {
type ParsedChunk,
parseChunks,
} from "@/components/assistant-ui/tool-ui-search-knowledge-base";
import { describe, expect, it } from "vitest";
// ── Tests ─────────────────────────────────────────────────────────────
describe("parseChunks — durable IDs (contracts §3/§4)", () => {
it("new XML with document_id + chunk_id populates all three identity fields", () => {
const xml = `
<chunk id="1" source="report.pdf" page="7" document_id="doc-abc" chunk_id="chunk-xyz">
The margin rose to 18%.
</chunk>
`.trim();
const parts: ParsedChunk[] = parseChunks(xml);
expect(parts).toHaveLength(1);
expect(parts[0].id).toBe("1");
expect(parts[0].documentId).toBe("doc-abc");
expect(parts[0].backendChunkId).toBe("chunk-xyz");
expect(parts[0].source).toBe("report.pdf");
expect(parts[0].page).toBe("7");
});
it("legacy XML without durable IDs leaves documentId and backendChunkId absent", () => {
// Old XML: no document_id, no chunk_id — hover-only, NOT preview-clickable (Q3).
const xml = `
<chunk id="2" source="old-doc.pdf" page="3">
Legacy chunk text.
</chunk>
`.trim();
const parts: ParsedChunk[] = parseChunks(xml);
expect(parts).toHaveLength(1);
expect(parts[0].id).toBe("2");
expect(parts[0].documentId).toBeUndefined();
expect(parts[0].backendChunkId).toBeUndefined();
});
it("citationId (id) is the visible counter string, never the backend UUID", () => {
const docId = "550e8400-e29b-41d4-a716-446655440000";
const chunkId = "6ba7b810-9dad-11d1-80b4-00c04fd430c8";
const xml = `<chunk id="5" source="paper.pdf" document_id="${docId}" chunk_id="${chunkId}">text</chunk>`;
const parts: ParsedChunk[] = parseChunks(xml);
expect(parts[0].id).toBe("5");
expect(parts[0].id).not.toBe(docId);
expect(parts[0].id).not.toBe(chunkId);
});
it("same filename in two chunks preserves distinct documentId values", () => {
const xml = `
<chunk id="1" source="annual.pdf" document_id="doc-001" chunk_id="chunk-001">First excerpt.</chunk>
<chunk id="2" source="annual.pdf" document_id="doc-002" chunk_id="chunk-002">Second excerpt.</chunk>
`.trim();
const parts: ParsedChunk[] = parseChunks(xml);
expect(parts).toHaveLength(2);
expect(parts[0].documentId).toBe("doc-001");
expect(parts[1].documentId).toBe("doc-002");
expect(parts[0].documentId).not.toBe(parts[1].documentId);
});
it("same visible id in different turns does not imply same backendChunkId", () => {
// Turn 1 and turn 2 both have id="1" but different backend identities.
const turn1 = `<chunk id="1" source="a.pdf" document_id="doc-A" chunk_id="chunk-A">Turn 1.</chunk>`;
const turn2 = `<chunk id="1" source="b.pdf" document_id="doc-B" chunk_id="chunk-B">Turn 2.</chunk>`;
const p1: ParsedChunk[] = parseChunks(turn1);
const p2: ParsedChunk[] = parseChunks(turn2);
expect(p1[0].id).toBe(p2[0].id); // both "1"
expect(p1[0].backendChunkId).not.toBe(p2[0].backendChunkId);
expect(p1[0].documentId).not.toBe(p2[0].documentId);
});
it("missing chunk_id only (partial durable attrs) → backendChunkId absent", () => {
const xml = `<chunk id="3" source="x.pdf" document_id="doc-XYZ">text</chunk>`;
const parts: ParsedChunk[] = parseChunks(xml);
expect(parts[0].documentId).toBe("doc-XYZ");
expect(parts[0].backendChunkId).toBeUndefined();
});
it("multiple new-format chunks all carry independent IDs", () => {
const xml = `
<chunk id="1" source="a.pdf" document_id="doc-1" chunk_id="ck-1">A</chunk>
<chunk id="2" source="b.pdf" document_id="doc-2" chunk_id="ck-2">B</chunk>
<chunk id="3" source="c.pdf" document_id="doc-3" chunk_id="ck-3">C</chunk>
`.trim();
const parts: ParsedChunk[] = parseChunks(xml);
expect(parts).toHaveLength(3);
const docIds = new Set(parts.map((p) => p.documentId));
const backendChunkIds = new Set(parts.map((p) => p.backendChunkId));
const citationIds = new Set(parts.map((p) => p.id));
expect(docIds.size).toBe(3);
expect(backendChunkIds.size).toBe(3);
expect(citationIds.size).toBe(3);
});
it("empty XML returns empty array without throwing", () => {
expect(parseChunks("")).toHaveLength(0);
expect(parseChunks("No chunks here.")).toHaveLength(0);
});
it("XML entity encoding in source attribute is decoded", () => {
// &amp; should decode to & in the source attribute (decodeXml in parseChunks)
const xml = `<chunk id="1" source="report &amp; summary.pdf" document_id="doc-1" chunk_id="ck-1">text</chunk>`;
const parts: ParsedChunk[] = parseChunks(xml);
expect(parts).toHaveLength(1);
expect(parts[0].id).toBe("1");
expect(parts[0].source).toBe("report & summary.pdf");
});
});

View file

@ -0,0 +1,92 @@
import type { RagDocument } from "@/features/rag/api/rag-api";
import { DocumentRow } from "@/features/rag/components/document-row";
import { fireEvent, render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import React from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
function makeDoc(overrides: Partial<RagDocument> = {}): RagDocument {
return {
id: "doc-abc",
kb_id: "kb-1",
thread_id: null,
filename: "report.pdf",
content_type: "application/pdf",
status: "completed",
num_chunks: 5,
byte_size: 10240,
error: null,
created_at: 1_700_000_000,
...overrides,
};
}
let onPreview: ReturnType<typeof vi.fn>;
let onDelete: ReturnType<typeof vi.fn>;
beforeEach(() => {
onPreview = vi.fn();
onDelete = vi.fn();
});
describe("DocumentRow preview event propagation", () => {
it("clicking a previewable row opens document-level preview", async () => {
render(
React.createElement(DocumentRow, {
doc: makeDoc(),
onPreview: onPreview as () => void,
onDelete: onDelete as () => void,
}),
);
await userEvent.click(
screen.getByRole("button", { name: /open preview of report.pdf/i }),
);
expect(onPreview).toHaveBeenCalledTimes(1);
});
it("Enter and Space open a previewable row", () => {
render(
React.createElement(DocumentRow, {
doc: makeDoc(),
onPreview: onPreview as () => void,
onDelete: onDelete as () => void,
}),
);
const row = screen.getByRole("button", {
name: /open preview of report.pdf/i,
});
fireEvent.keyDown(row, { key: "Enter" });
fireEvent.keyDown(row, { key: " " });
expect(onPreview).toHaveBeenCalledTimes(2);
});
it("clicking delete does not open preview", async () => {
render(
React.createElement(DocumentRow, {
doc: makeDoc(),
onPreview: onPreview as () => void,
onDelete: onDelete as () => void,
}),
);
await userEvent.click(screen.getByRole("button", { name: /delete/i }));
expect(onDelete).toHaveBeenCalledTimes(1);
expect(onPreview).not.toHaveBeenCalled();
});
it("non-previewable rows have no row button semantics", () => {
render(
React.createElement(DocumentRow, {
doc: makeDoc({ status: "pending" }),
onDelete: onDelete as () => void,
}),
);
expect(screen.queryByRole("button", { name: /open preview/i })).toBeNull();
});
});

View file

@ -0,0 +1,77 @@
import { KnowledgeBasesTab } from "@/features/settings/tabs/knowledge-bases-tab";
import { render, screen } from "@testing-library/react";
import React from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
const { mockUsePreviewStore } = vi.hoisted(() => {
const state = {
target: null as unknown,
status: "idle",
close: vi.fn(),
};
const fn = vi.fn((selector?: (s: typeof state) => unknown) => {
if (typeof selector === "function") return selector(state);
return state;
}) as ReturnType<typeof vi.fn> & { __state: typeof state };
fn.__state = state;
return { mockUsePreviewStore: fn };
});
vi.mock("@/features/rag/stores/preview-store", () => ({
usePreviewStore: mockUsePreviewStore,
}));
vi.mock("@/features/rag/components/kb-list", () => ({
KBList: () => React.createElement("div", { "data-testid": "kb-list" }),
}));
vi.mock("@/features/rag/components/kb-detail-panel", () => ({
KBDetailPanel: () =>
React.createElement("div", { "data-testid": "kb-detail-panel" }),
}));
vi.mock("@/features/rag/components/preview-panel", () => ({
PreviewPanel: ({ open }: { open: boolean }) =>
React.createElement("div", {
"data-testid": "settings-preview-panel",
"data-open": String(open),
}),
}));
vi.mock("@/features/rag/components/thread-index-list", () => ({
ThreadIndexList: () =>
React.createElement("div", { "data-testid": "thread-index-list" }),
}));
vi.mock("@/features/rag/components/rag-defaults-section", () => ({
RagDefaultsSection: () =>
React.createElement("div", { "data-testid": "rag-defaults-section" }),
}));
beforeEach(() => {
mockUsePreviewStore.__state.target = null;
mockUsePreviewStore.__state.status = "idle";
mockUsePreviewStore.mockImplementation(
(selector?: (s: typeof mockUsePreviewStore.__state) => unknown) => {
if (typeof selector === "function") {
return selector(mockUsePreviewStore.__state);
}
return mockUsePreviewStore.__state;
},
);
});
describe("KnowledgeBasesTab preview host", () => {
it("renders a preview panel when the preview store is active", () => {
mockUsePreviewStore.__state.target = {
documentId: "doc-abc",
filename: "report.pdf",
};
mockUsePreviewStore.__state.status = "ready";
render(React.createElement(KnowledgeBasesTab));
const panel = screen.getByTestId("settings-preview-panel");
expect(panel.getAttribute("data-open")).toBe("true");
});
});

View file

@ -0,0 +1,93 @@
import type { PreviewTarget } from "@/features/rag/api/rag-api";
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
import React from "react";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
const { mockFetchPreviewTarget, mockFetchPreviewFileBlob } = vi.hoisted(() => ({
mockFetchPreviewTarget:
vi.fn<
(documentId: string, chunkId?: string | null) => Promise<PreviewTarget>
>(),
mockFetchPreviewFileBlob:
vi.fn<(documentId: string, signal?: AbortSignal) => Promise<Blob>>(),
}));
vi.mock("@/features/rag/api/rag-api", async (importOriginal) => {
const original =
await importOriginal<typeof import("@/features/rag/api/rag-api")>();
return {
...original,
fetchPreviewTarget: mockFetchPreviewTarget,
fetchPreviewFileBlob: mockFetchPreviewFileBlob,
};
});
vi.mock("react-pdf", () => ({
Document: ({ children }: { children: React.ReactNode }) =>
React.createElement("div", { "data-testid": "pdf-document" }, children),
Page: () => React.createElement("div", { "data-testid": "pdf-page" }),
pdfjs: { GlobalWorkerOptions: { workerSrc: "" } },
}));
import { PreviewPanel } from "@/features/rag/components/preview-panel";
import { usePreviewStore } from "@/features/rag/stores/preview-store";
function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
return {
documentId: "doc-abc",
filename: "report.html",
contentType: "text/plain",
mediaKind: "html",
byteSize: 100,
status: "completed",
kbId: "kb-1",
threadId: null,
chunkId: "chunk-1",
chunkIndex: 0,
targetPage: 1,
snippet: "safe extracted text",
kind: "text",
imageUrl: null,
sourcePageIndex: null,
pageCharStart: null,
pageCharEnd: null,
lineStart: null,
lineEnd: null,
pdfRegions: [],
...overrides,
};
}
beforeEach(() => {
mockFetchPreviewTarget.mockReset();
mockFetchPreviewFileBlob.mockReset();
usePreviewStore.getState().close();
});
afterEach(() => {
usePreviewStore.getState().close();
});
describe("preview a11y hardening", () => {
it("Escape closes the preview and restores focus to the opener", async () => {
const opener = document.createElement("button");
opener.textContent = "Open preview";
document.body.appendChild(opener);
opener.focus();
mockFetchPreviewTarget.mockResolvedValue(target());
await usePreviewStore.getState().open({ documentId: "doc-abc" });
render(React.createElement(PreviewPanel, { open: true }));
expect(
screen.getByRole("region", { name: /document preview/i }),
).toBeInTheDocument();
fireEvent.keyDown(document, { key: "Escape" });
await waitFor(() => {
expect(usePreviewStore.getState().status).toBe("idle");
});
expect(document.activeElement).toBe(opener);
opener.remove();
});
});

View file

@ -0,0 +1,521 @@
/**
* Tests for preview-panel HTML/DOCX/unknown must NEVER render inline (T5 / Risk #3).
*
* Acceptance criteria (contracts §5.4, PLAN.md T5, decisions Q7):
* - mediaKind === "pdf" react-pdf view is mounted (or loading indicator shown).
* - mediaKind === "html" text-view fallback shown, NO object/embed/iframe with blob URL.
* - mediaKind === "docx" text-view fallback shown, NO inline rendering.
* - mediaKind === "unknown" unavailable/download state, NOT inline.
* - mediaKind === "text" text/snippet view shown.
* - Panel without a target renders nothing or unavailable state.
*/
import {
type PreviewMediaKind,
type PreviewTarget,
} from "@/features/rag/api/rag-api";
import type { PreviewLoadStatus } from "@/features/rag/stores/preview-store";
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
import React from "react";
import {
type MockInstance,
afterEach,
beforeEach,
describe,
expect,
it,
vi,
} from "vitest";
const DOWNLOAD_BUTTON_NAME = /download/i;
const LONG_CONTENT_TEXT = /Some long content/;
// ── Mock preview store ────────────────────────────────────────────────
// The real component uses per-field selectors: usePreviewStore((s) => s.target)
// so the mock must handle the selector pattern.
// vi.hoisted ensures the mock fn is initialised before vi.mock factory runs.
interface MockStoreState {
target: PreviewTarget | null;
previewBlobUrl: string | null;
previewBlob: Blob | null;
previewFileUrl: string | null;
previewFileUrlExpiresAt: number | null;
status: PreviewLoadStatus;
error: string | null;
close: () => void;
open: () => void;
}
let mockState: MockStoreState = {
target: null,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "idle",
error: null,
close: vi.fn(),
open: vi.fn(),
};
const { mockAuthFetch, mockUsePreviewStore } = vi.hoisted(() => {
// usePreviewStore is called two ways:
// usePreviewStore((s) => s.field) — selector form (React hook)
// usePreviewStore.getState().close() — outside React (cleanup effect)
const fn = vi.fn((selector?: (s: MockStoreState) => unknown) => {
if (typeof selector === "function") {
return selector(mockState);
}
return mockState;
}) as ReturnType<typeof vi.fn> & { getState: () => MockStoreState };
fn.getState = () => mockState;
return { mockAuthFetch: vi.fn(), mockUsePreviewStore: fn };
});
vi.mock("@/features/auth", () => ({
authFetch: mockAuthFetch,
getAuthToken: () => "mock-token-123",
}));
vi.mock("@/features/rag/stores/preview-store", async (importOriginal) => {
const real =
await importOriginal<
typeof import("@/features/rag/stores/preview-store")
>();
return {
...real,
usePreviewStore: mockUsePreviewStore,
// isInlineBlobAllowed passes through from the real module so assertions
// use the production allowlist, not a test-local copy (D1.5 fix).
};
});
// react-pdf requires a browser worker URL that doesn't exist in jsdom.
vi.mock("react-pdf", () => ({
Document: ({ children }: { children: React.ReactNode }) =>
React.createElement("div", { "data-testid": "pdf-document" }, children),
Page: () => React.createElement("div", { "data-testid": "pdf-page" }),
pdfjs: { GlobalWorkerOptions: { workerSrc: "" } },
}));
beforeEach(() => {
mockAuthFetch.mockReset();
mockAuthFetch.mockResolvedValue(
new Response(new Blob(["download bytes"], { type: "text/plain" }), {
status: 200,
}),
);
mockState = {
target: null,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "idle",
error: null,
close: vi.fn(),
open: vi.fn(),
};
mockUsePreviewStore.mockImplementation(
(selector?: (s: MockStoreState) => unknown) => {
if (typeof selector === "function") {
return selector(mockState);
}
return mockState;
},
);
// Restore getState after mockImplementation replaces the fn internals
mockUsePreviewStore.getState = () => mockState;
// Mock window.matchMedia globally for tests
window.matchMedia = vi.fn().mockImplementation((query) => ({
matches: false,
media: query,
onchange: null,
addListener: vi.fn(),
removeListener: vi.fn(),
addEventListener: vi.fn(),
removeEventListener: vi.fn(),
dispatchEvent: vi.fn(),
}));
});
// ── Import panel + production allowlist predicate AFTER mocks ─────────
import { PreviewPanel } from "@/features/rag/components/preview-panel";
import { isInlineBlobAllowed } from "@/features/rag/stores/preview-store";
// ── Helpers ───────────────────────────────────────────────────────────
function makeTarget(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
return {
documentId: "doc-abc",
filename: "report.pdf",
contentType: "application/pdf",
mediaKind: "pdf",
byteSize: 100,
status: "completed",
kbId: "kb-1",
threadId: null,
chunkId: null,
chunkIndex: null,
targetPage: null,
snippet: null,
kind: null,
imageUrl: null,
sourcePageIndex: null,
pageCharStart: null,
pageCharEnd: null,
lineStart: null,
lineEnd: null,
pdfRegions: [],
...overrides,
};
}
function setReady(overrides: Partial<PreviewTarget> = {}): void {
mockState.target = makeTarget(overrides);
mockState.status = "ready";
}
// ── Tests against real panel component ───────────────────────────────
describe("preview-panel inline rendering safety (contracts §5.4 / Risk #3)", () => {
it("panel with no target (idle) renders without crashing", () => {
const { container } = render(
React.createElement(PreviewPanel, { open: true }),
);
expect(container).toBeDefined();
expect(document.querySelector("iframe")).toBeNull();
});
it("html mediaKind does not render an iframe, object, or embed element", () => {
setReady({ mediaKind: "html", filename: "malicious.html" });
render(React.createElement(PreviewPanel, { open: true }));
expect(document.querySelector("iframe")).toBeNull();
expect(document.querySelector("object")).toBeNull();
expect(document.querySelector("embed")).toBeNull();
});
it("docx mediaKind does not render an iframe, object, or embed element", () => {
setReady({ mediaKind: "docx", filename: "report.docx" });
render(React.createElement(PreviewPanel, { open: true }));
expect(document.querySelector("iframe")).toBeNull();
expect(document.querySelector("object")).toBeNull();
expect(document.querySelector("embed")).toBeNull();
});
it("unknown mediaKind does not render inline blob content", () => {
setReady({ mediaKind: "unknown", filename: "data.bin" });
render(React.createElement(PreviewPanel, { open: true }));
expect(document.querySelector("iframe")).toBeNull();
expect(document.querySelector("object")).toBeNull();
expect(document.querySelector("embed")).toBeNull();
});
it.each<PreviewMediaKind>(["html", "docx", "unknown"])(
"%s download creates only a download object URL, never inline preview content",
async (mediaKind) => {
const createObjectUrl = vi
.spyOn(URL, "createObjectURL")
.mockReturnValue("blob:unsafe");
const revokeObjectUrl = vi
.spyOn(URL, "revokeObjectURL")
.mockImplementation(() => undefined);
setReady({
mediaKind,
filename: `unsafe.${mediaKind}`,
contentType: "text/plain",
snippet: "Extracted text only.",
});
render(React.createElement(PreviewPanel, { open: true }));
fireEvent.click(
screen.getByRole("button", { name: DOWNLOAD_BUTTON_NAME }),
);
await waitFor(() => {
expect(mockAuthFetch).toHaveBeenCalledWith(
"/api/rag/documents/doc-abc/file",
);
});
expect(createObjectUrl).toHaveBeenCalledWith(expect.any(Blob));
await waitFor(() => {
expect(revokeObjectUrl).toHaveBeenCalledWith("blob:unsafe");
});
expect(document.querySelector("iframe")).toBeNull();
expect(document.querySelector("object")).toBeNull();
expect(document.querySelector("embed")).toBeNull();
createObjectUrl.mockRestore();
revokeObjectUrl.mockRestore();
},
);
it("text mediaKind renders without iframe (text fallback path)", () => {
setReady({
mediaKind: "text",
filename: "notes.txt",
snippet: "This is the extracted text content.",
});
render(React.createElement(PreviewPanel, { open: true }));
expect(document.querySelector("iframe")).toBeNull();
});
it("open=false triggers close side-effect on the store", () => {
const closeFn = vi.fn();
mockState.close = closeFn;
mockState.target = makeTarget();
mockState.status = "ready";
const { rerender } = render(
React.createElement(PreviewPanel, { open: true }),
);
rerender(React.createElement(PreviewPanel, { open: false }));
// The useEffect for open=false should have called close()
expect(closeFn).toHaveBeenCalled();
});
it("closes the preview panel on Escape key down (Escape key closures)", () => {
const closeFn = vi.fn();
mockState.close = closeFn;
mockState.target = makeTarget();
mockState.status = "ready";
render(React.createElement(PreviewPanel, { open: true }));
fireEvent.keyDown(document, { key: "Escape" });
expect(closeFn).toHaveBeenCalled();
});
it("renders with premium glassmorphic visual details and a pulsing green indicator dot", () => {
setReady({ mediaKind: "text", filename: "notes.txt" });
render(React.createElement(PreviewPanel, { open: true }));
const section = screen.getByLabelText("Document preview");
expect(section).toHaveClass("bg-panel-surface/85");
expect(section).toHaveClass("backdrop-blur-lg");
expect(section).toHaveClass("border-border/40");
expect(section).toHaveClass("shadow-lg");
// Pulser dot
const pulser = section.querySelector(".animate-pulse");
expect(pulser).toBeInTheDocument();
expect(pulser).toHaveClass("bg-primary");
expect(pulser).toHaveClass("w-2");
expect(pulser).toHaveClass("h-2");
});
it("shifts the layout to a full mobile Sheet drawer overlay when the viewport is squeezed (< 1024px)", () => {
window.matchMedia = vi.fn().mockImplementation((query) => ({
matches: true,
media: query,
onchange: null,
addListener: vi.fn(),
removeListener: vi.fn(),
addEventListener: vi.fn(),
removeEventListener: vi.fn(),
dispatchEvent: vi.fn(),
}));
setReady({ mediaKind: "text", filename: "notes.txt" });
render(React.createElement(PreviewPanel, { open: true }));
// Radix UI Sheet component should render dialog role in mobile viewports
const dialog = screen.getByRole("dialog");
expect(dialog).toBeInTheDocument();
expect(dialog).toHaveClass("preview-sheet-content");
expect(screen.getByText("Document preview")).toBeInTheDocument();
});
});
// ── Pure-logic: inline allowlist (always green) ───────────────────────
// Uses the real production isInlineBlobAllowed (D1.5 fix: no local copy).
describe("inline object URL allowlist (contracts §5.4, pure logic)", () => {
const inlineSafe: PreviewMediaKind[] = ["pdf", "text", "image"];
const inlineUnsafe: PreviewMediaKind[] = ["html", "docx", "unknown"];
it.each(inlineSafe)("mediaKind=%s is inline-safe", (mk) => {
expect(isInlineBlobAllowed(mk)).toBe(true);
});
it.each(inlineUnsafe)("mediaKind=%s is NOT inline-safe (Risk #3)", (mk) => {
expect(isInlineBlobAllowed(mk)).toBe(false);
});
});
describe("preview-panel stable scrollbars, sheets, layouts, and downloads", () => {
let createObjectUrl: MockInstance<typeof URL.createObjectURL>;
let revokeObjectUrl: MockInstance<typeof URL.revokeObjectURL>;
beforeEach(() => {
createObjectUrl = vi
.spyOn(URL, "createObjectURL")
.mockReturnValue("blob:safe-url");
revokeObjectUrl = vi
.spyOn(URL, "revokeObjectURL")
.mockImplementation(() => undefined);
});
afterEach(() => {
createObjectUrl.mockRestore();
revokeObjectUrl.mockRestore();
});
it("asserts stable scrollbar style classes are present on panel content", () => {
setReady({
mediaKind: "text",
filename: "notes.txt",
snippet: "Some long content that requires scrolling ".repeat(20),
});
render(React.createElement(PreviewPanel, { open: true }));
// The snippet is rendered in a <pre> element. Check if it has overflow-auto
const preElement = screen.getByText(LONG_CONTENT_TEXT);
expect(preElement).toHaveClass("overflow-auto");
expect(preElement).toHaveClass("flex-1");
});
it("asserts non-nested sheets are rendered in squeezed viewports", () => {
window.matchMedia = vi.fn().mockImplementation((query) => ({
matches: true,
media: query,
onchange: null,
addListener: vi.fn(),
removeListener: vi.fn(),
addEventListener: vi.fn(),
removeEventListener: vi.fn(),
dispatchEvent: vi.fn(),
}));
setReady({ mediaKind: "text", filename: "notes.txt" });
render(React.createElement(PreviewPanel, { open: true }));
const dialogs = screen.getAllByRole("dialog");
expect(dialogs.length).toBe(1);
expect(dialogs[0]).toHaveClass("preview-sheet-content");
const nestedDialogs = dialogs[0].querySelectorAll("[role='dialog']");
expect(nestedDialogs.length).toBe(0);
});
it("supports responsive collapses under different viewport widths", () => {
const mockMatchMedia = vi.fn().mockImplementation((query) => ({
matches: query.includes("max-width: 1023px"),
media: query,
onchange: null,
addListener: vi.fn(),
removeListener: vi.fn(),
addEventListener: vi.fn(),
removeEventListener: vi.fn(),
dispatchEvent: vi.fn(),
}));
window.matchMedia = mockMatchMedia;
setReady({ mediaKind: "text", filename: "notes.txt" });
const { unmount } = render(
React.createElement(PreviewPanel, { open: true }),
);
expect(screen.getByRole("dialog")).toBeInTheDocument();
unmount();
window.matchMedia = vi.fn().mockImplementation((query) => ({
matches: false,
media: query,
onchange: null,
addListener: vi.fn(),
removeListener: vi.fn(),
addEventListener: vi.fn(),
removeEventListener: vi.fn(),
dispatchEvent: vi.fn(),
}));
render(React.createElement(PreviewPanel, { open: true }));
expect(screen.queryByRole("dialog")).toBeNull();
expect(screen.getByLabelText("Document preview")).toBeInTheDocument();
});
it("asserts object URL download logic: uses URL.createObjectURL for safe types and Data URL for unsafe types", async () => {
setReady({
mediaKind: "text",
filename: "notes.txt",
contentType: "text/plain",
snippet: "Text snippet.",
});
render(React.createElement(PreviewPanel, { open: true }));
fireEvent.click(screen.getByRole("button", { name: DOWNLOAD_BUTTON_NAME }));
await waitFor(() => {
expect(mockAuthFetch).toHaveBeenCalledWith(
"/api/rag/documents/doc-abc/file",
);
});
expect(createObjectUrl).toHaveBeenCalled();
});
});
describe("PreviewTextView precise highlights matching", () => {
it("highlights with character ranges", () => {
setReady({
mediaKind: "text",
filename: "notes.txt",
snippet: "Line 1: Hello World\nLine 2: Target Phrase\nLine 3: Goodbye",
pageCharStart: 28,
pageCharEnd: 41,
});
render(React.createElement(PreviewPanel, { open: true }));
const mark = screen.getByText("Target Phrase");
expect(mark.tagName).toBe("MARK");
expect(mark).toHaveClass("bg-primary/20", "ring-primary/60");
});
it("highlights with line numbers", () => {
setReady({
mediaKind: "text",
filename: "notes.txt",
snippet: "Line one text\nLine two text\nLine three text",
lineStart: 2,
lineEnd: 2,
});
render(React.createElement(PreviewPanel, { open: true }));
const mark = screen.getByText("Line two text");
expect(mark.tagName).toBe("MARK");
expect(mark).toHaveClass("bg-primary/20", "ring-primary/60");
});
it("highlights with fuzzy fallback matching high density line", () => {
setReady({
mediaKind: "text",
filename: "notes.txt",
snippet: "...\nAlphanumericDensity123456\n...",
lineStart: 999, // Trigger hasLocator without matching any specific line range
});
render(React.createElement(PreviewPanel, { open: true }));
const mark = screen.getByText("AlphanumericDensity123456");
expect(mark.tagName).toBe("MARK");
});
});

View file

@ -0,0 +1,386 @@
import type { PreviewTarget } from "@/features/rag/api/rag-api";
import type { PreviewPdfRegion } from "@/features/rag/api/rag-api";
import { PreviewPdfView } from "@/features/rag/components/preview-pdf-view";
import {
act,
fireEvent,
render,
screen,
waitFor,
within,
} from "@testing-library/react";
import React from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
const SOURCE_EXCERPT_TEXT = /source excerpt/i;
/** The thumbnail rail also renders mocked `<Page>` elements, so every
* test that targets the main render must scope through the
* `pdf-main-page` wrapper instead of taking the first `pdf-page`. */
async function findMainPdfPage(): Promise<HTMLElement> {
const wrapper = await screen.findByTestId("pdf-main-page");
return within(wrapper).getByTestId("pdf-page");
}
function getMainPdfPage(): HTMLElement {
const wrapper = screen.getByTestId("pdf-main-page");
return within(wrapper).getByTestId("pdf-page");
}
vi.mock("react-pdf", () => ({
Document: ({
children,
file,
onLoadSuccess,
}: {
children: React.ReactNode;
file?: unknown;
onLoadSuccess?: (result: { numPages: number }) => void;
}) => {
onLoadSuccess?.({ numPages: 1 });
return React.createElement(
"div",
{
"data-testid": "pdf-document",
"data-file-kind": file instanceof Blob ? "blob" : typeof file,
"data-file-url":
file && typeof file === "object" && "url" in file
? String((file as { url: string }).url)
: "",
},
children,
);
},
Page: ({
customTextRenderer,
width,
renderTextLayer,
}: {
customTextRenderer?: (item: { str: string }) => string;
width?: number;
renderTextLayer?: boolean;
}) => {
const html =
customTextRenderer?.({ str: "target phrase" }) ?? "target phrase";
return React.createElement("div", {
"data-testid": "pdf-page",
"data-width": String(width ?? ""),
"data-render-text-layer": String(renderTextLayer),
"data-rendered-html": html,
});
},
pdfjs: { GlobalWorkerOptions: { workerSrc: "" } },
}));
function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
return {
documentId: "doc-abc",
filename: "report.pdf",
contentType: "application/pdf",
mediaKind: "pdf",
byteSize: 100,
status: "completed",
kbId: "kb-1",
threadId: null,
chunkId: "chunk-1",
chunkIndex: 0,
targetPage: 1,
snippet: "target phrase appears here",
kind: "text",
imageUrl: null,
sourcePageIndex: 0,
pageCharStart: 0,
pageCharEnd: 13,
lineStart: 1,
lineEnd: 1,
pdfRegions: [],
...overrides,
};
}
beforeEach(() => {
class ResizeObserverMock implements ResizeObserver {
observe(_target: Element, _options?: ResizeObserverOptions) {
// jsdom has no layout observer; the component only needs the API shape.
}
unobserve(_target: Element) {
// jsdom has no layout observer; the component only needs the API shape.
}
disconnect() {
// jsdom has no layout observer; the component only needs the API shape.
}
}
vi.stubGlobal("ResizeObserver", ResizeObserverMock);
});
describe("PreviewPdfView smoke", () => {
it("renders a range URL source with text search and exact region overlay", async () => {
render(
React.createElement(PreviewPdfView, {
target: target({
pdfRegions: [
{
pageIndex: 0,
pageNumber: 1,
x: 0.1,
y: 0.2,
width: 0.3,
height: 0.04,
confidence: "exact",
source: "pymupdf-search",
},
],
}),
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
}),
);
expect(screen.getByTestId("pdf-document")).toHaveAttribute(
"data-file-kind",
"object",
);
expect(screen.getByTestId("pdf-document")).toHaveAttribute(
"data-file-url",
expect.stringContaining("/file-signed?token=signed"),
);
const page = await findMainPdfPage();
await waitFor(() => {
expect(page.getAttribute("data-render-text-layer")).toBe("true");
});
expect(page.getAttribute("data-rendered-html")).toBe("target phrase");
expect(page.getAttribute("data-rendered-html")).not.toContain("<mark>");
// Verify brand green highlight overlays
const regionHighlight = screen.getByTestId("pdf-region-highlight");
expect(regionHighlight).toBeInTheDocument();
expect(regionHighlight).toHaveClass("bg-primary/20");
expect(regionHighlight).toHaveClass("ring-primary/60");
// Verify Tailwind v4 light-mode isolation reset wrapper. After the
// thumbnail-rail refactor, the light wrapper lives INSIDE the
// Document and directly wraps the main-page block.
const wrapper = screen.getByTestId("pdf-main-page").parentElement;
expect(wrapper).toHaveClass("light");
expect(wrapper).toHaveClass("bg-white");
expect(wrapper).toHaveClass("text-slate-900");
// Verify Shadcn toolbar elements and rounded-full pill groups
const zoomInBtn = screen.getByRole("button", { name: "Zoom in" });
expect(zoomInBtn).toHaveClass("rounded-full");
expect(zoomInBtn.parentElement).toHaveClass(
"bg-muted/40",
"p-0.5",
"shadow-xs",
);
// Source-excerpt card uses a neutral muted surface (no brand-coloured
// left rail) so it sits inside the panel without visually competing.
const excerptCard = screen.getByText(SOURCE_EXCERPT_TEXT).parentElement;
expect(excerptCard).toHaveClass("border-border/60");
expect(excerptCard).toHaveClass("bg-muted/30");
expect(excerptCard).not.toHaveClass("border-l-primary");
fireEvent.change(screen.getByLabelText("Search this PDF"), {
target: { value: "phrase" },
});
await waitFor(() => {
expect(
getMainPdfPage().getAttribute("data-rendered-html"),
).toContain("<mark>phrase</mark>");
});
const beforeZoom = Number(page.getAttribute("data-width"));
fireEvent.click(screen.getByRole("button", { name: "Zoom in" }));
await waitFor(() => {
expect(
Number(getMainPdfPage().getAttribute("data-width")),
).toBeGreaterThan(beforeZoom);
});
});
it("debounces ResizeObserver transitions to prevent infinite rendering loops", async () => {
const resizeCallbacks: ResizeObserverCallback[] = [];
class FakeResizeObserver implements ResizeObserver {
constructor(callback: ResizeObserverCallback) {
resizeCallbacks.push(callback);
}
observe(_target: Element, _options?: ResizeObserverOptions) {
// jsdom has no layout observer; the component only needs the API shape.
}
unobserve(_target: Element) {
// jsdom has no layout observer; the component only needs the API shape.
}
disconnect() {
// jsdom has no layout observer; the component only needs the API shape.
}
}
vi.stubGlobal("ResizeObserver", FakeResizeObserver);
render(
React.createElement(PreviewPdfView, {
target: target(),
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
}),
);
// Initial render sets width synchronously on mount. Let's capture the initial width.
const page = await findMainPdfPage();
const initialWidth = Number(page.getAttribute("data-width"));
// Activate fake timers AFTER finding the elements to avoid findByTestId timeout
vi.useFakeTimers();
// Set up HTMLDivElement.prototype.clientWidth mock
const originalClientWidth = Object.getOwnPropertyDescriptor(
HTMLDivElement.prototype,
"clientWidth",
);
let clientWidthValue = 300;
Object.defineProperty(HTMLDivElement.prototype, "clientWidth", {
get() {
return clientWidthValue;
},
configurable: true,
});
// Now trigger resize callback after changing clientWidth
clientWidthValue = 600;
const resizeCallback = resizeCallbacks[0];
if (!resizeCallback) {
throw new Error("Expected ResizeObserver callback to be registered");
}
const resizeObserver: ResizeObserver = {
observe() {
// The callback under test ignores the observer instance.
},
unobserve() {
// The callback under test ignores the observer instance.
},
disconnect() {
// The callback under test ignores the observer instance.
},
};
resizeCallback([], resizeObserver);
// Width should NOT be updated immediately because of the 100ms debounce
expect(Number(getMainPdfPage().getAttribute("data-width"))).toBe(
initialWidth,
);
// Fast-forward time by 100ms to trigger the debounced callback and flush updates
act(() => {
vi.advanceTimersByTime(100);
vi.runAllTimers();
});
// Now the width should have updated
expect(
Number(getMainPdfPage().getAttribute("data-width")),
).not.toBe(initialWidth);
expect(Number(getMainPdfPage().getAttribute("data-width"))).toBe(572); // 600 - 28 (PDF_BODY_GUTTER_PX)
// Clean up prototype descriptor
if (originalClientWidth) {
Object.defineProperty(
HTMLDivElement.prototype,
"clientWidth",
originalClientWidth,
);
} else {
Reflect.deleteProperty(HTMLDivElement.prototype, "clientWidth");
}
vi.useRealTimers();
});
it("renders only 'exact' confidence highlights and positions them with correct percentages", async () => {
const nonExactRegion = {
pageIndex: 0,
pageNumber: 1,
x: 0.5,
y: 0.5,
width: 0.2,
height: 0.2,
confidence: "fuzzy",
source: "pymupdf-search",
} as unknown as PreviewPdfRegion;
render(
React.createElement(PreviewPdfView, {
target: target({
pdfRegions: [
{
pageIndex: 0,
pageNumber: 1,
x: 0.15,
y: 0.25,
width: 0.35,
height: 0.45,
confidence: "exact",
source: "pymupdf-search",
},
nonExactRegion,
],
}),
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
}),
);
await findMainPdfPage();
const highlights = screen.getAllByTestId("pdf-region-highlight");
expect(highlights.length).toBe(1);
const exactHighlight = highlights[0];
expect(exactHighlight.style.left).toBe("15%");
expect(exactHighlight.style.top).toBe("25%");
expect(exactHighlight.style.width).toBe("35%");
expect(exactHighlight.style.height).toBe("45%");
});
it("uses stable scrollbar style classes in the PDF sidebar and page container to prevent shifting", async () => {
render(
React.createElement(PreviewPdfView, {
target: target(),
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
}),
);
const mainPageWrapper = await screen.findByTestId("pdf-main-page");
// pdf-main-page → light-wrapper → scrollContainer
const scrollContainer =
mainPageWrapper.parentElement?.parentElement ?? null;
expect(scrollContainer).toHaveClass("preview-scrollbar");
expect(scrollContainer).toHaveClass("overflow-y-scroll");
expect(scrollContainer).toHaveClass("overflow-x-auto");
const sidebar = screen.getByRole("button", {
name: "Go to page 1",
}).parentElement;
expect(sidebar).toHaveClass("preview-scrollbar");
expect(sidebar).toHaveClass("overflow-y-auto");
});
it("highlights search terms using the custom text renderer with the mark wrapper", async () => {
render(
React.createElement(PreviewPdfView, {
target: target({
snippet: "this snippet contains some special keyword",
}),
file: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
}),
);
await findMainPdfPage();
fireEvent.change(screen.getByLabelText("Search this PDF"), {
target: { value: "phrase" },
});
await waitFor(() => {
expect(
getMainPdfPage().getAttribute("data-rendered-html"),
).toContain("<mark>phrase</mark>");
});
});
});

View file

@ -0,0 +1,285 @@
/**
* Tests for preview-store object URL lifecycle (contracts §5, T4).
*
* Coverage:
* - open() revokes previous object URL before assigning a new one.
* - close() revokes any live object URL.
* - Opening doc B while doc A is loaded revokes doc A's URL.
* - PDFs use a signed range URL instead of a full blob download.
* - Inline object URLs are created ONLY for safe non-PDF mediaKind (text/image).
* - For unsafe mediaKind (html/docx/unknown) blob fetch is skipped; previewBlobUrl = null.
* - isInlineBlobAllowed pure predicate matches contracts §5.4 allowlist.
* - __previewStoreInternals() verifies module-scoped cleanup.
*/
import type {
PreviewMediaKind,
PreviewTarget,
} from "@/features/rag/api/rag-api";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
// ── Mock rag-api BEFORE importing the store ───────────────────────────
// vi.hoisted ensures the mock refs are initialised before vi.mock factory
// runs (vi.mock is hoisted to the top of the file by Vitest's transformer).
const {
mockFetchPreviewTarget,
mockFetchPreviewFileBlob,
mockFetchPreviewFileUrl,
} = vi.hoisted(() => ({
mockFetchPreviewTarget:
vi.fn<
(documentId: string, chunkId?: string | null) => Promise<PreviewTarget>
>(),
mockFetchPreviewFileBlob:
vi.fn<(documentId: string, signal?: AbortSignal) => Promise<Blob>>(),
mockFetchPreviewFileUrl: vi.fn<
(
documentId: string,
signal?: AbortSignal,
) => Promise<{ url: string; expiresAt: number }>
>(),
}));
vi.mock("@/features/rag/api/rag-api", async (importOriginal) => {
const original =
await importOriginal<typeof import("@/features/rag/api/rag-api")>();
return {
...original,
fetchPreviewTarget: mockFetchPreviewTarget,
fetchPreviewFileBlob: mockFetchPreviewFileBlob,
fetchPreviewFileUrl: mockFetchPreviewFileUrl,
};
});
// ── Import store AFTER mock registration ─────────────────────────────
import {
__previewStoreInternals,
isInlineBlobAllowed,
usePreviewStore,
} from "@/features/rag/stores/preview-store";
// ── Mock URL.createObjectURL / revokeObjectURL ────────────────────────
let urlCounter = 0;
beforeEach(() => {
urlCounter = 0;
vi.spyOn(URL, "createObjectURL").mockImplementation(() => {
return `blob:test/${++urlCounter}`;
});
vi.spyOn(URL, "revokeObjectURL").mockImplementation((_url: string) => {
/* no-op */
});
mockFetchPreviewTarget.mockReset();
mockFetchPreviewFileBlob.mockReset();
mockFetchPreviewFileUrl.mockReset();
// Reset store to idle between tests
usePreviewStore.getState().close();
});
afterEach(() => {
vi.restoreAllMocks();
});
// ── Helpers ───────────────────────────────────────────────────────────
function makeTarget(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
return {
documentId: "doc-abc",
filename: "report.pdf",
contentType: "application/pdf",
mediaKind: "pdf",
byteSize: 100,
status: "completed",
kbId: "kb-1",
threadId: null,
chunkId: null,
chunkIndex: null,
targetPage: null,
snippet: null,
kind: null,
imageUrl: null,
sourcePageIndex: null,
pageCharStart: null,
pageCharEnd: null,
lineStart: null,
lineEnd: null,
pdfRegions: [],
...overrides,
};
}
function makePdfBlob(): Blob {
return new Blob(["%PDF-1.4"], { type: "application/pdf" });
}
// ── Pure-logic: isInlineBlobAllowed (always green) ────────────────────
describe("isInlineBlobAllowed (contracts §5.4, pure logic)", () => {
const safe: PreviewMediaKind[] = ["pdf", "text", "image"];
const unsafe: PreviewMediaKind[] = ["html", "docx", "unknown"];
it.each(safe)("mediaKind=%s is inline-safe", (mk) => {
expect(isInlineBlobAllowed(mk)).toBe(true);
});
it.each(unsafe)("mediaKind=%s is NOT inline-safe (Risk #3)", (mk) => {
expect(isInlineBlobAllowed(mk)).toBe(false);
});
});
// ── Integration tests against real store ─────────────────────────────
describe("preview-store open/close lifecycle (contracts §5)", () => {
it("open() for pdf stores a signed URL without creating an object URL", async () => {
mockFetchPreviewTarget.mockResolvedValue(makeTarget({ mediaKind: "pdf" }));
mockFetchPreviewFileUrl.mockResolvedValue({
url: "http://127.0.0.1:8888/api/rag/documents/doc-abc/file-signed?token=signed",
expiresAt: 1_700_000_000,
});
await usePreviewStore.getState().open({ documentId: "doc-abc" });
expect(mockFetchPreviewFileBlob).not.toHaveBeenCalled();
expect(URL.createObjectURL).not.toHaveBeenCalled();
const { previewBlob, previewBlobUrl, previewFileUrl, status } =
usePreviewStore.getState();
expect(previewBlob).toBeNull();
expect(previewBlobUrl).toBeNull();
expect(previewFileUrl).toContain("/file-signed?token=signed");
expect(status).toBe("ready");
});
it("open() passes backendChunkId through to fetchPreviewTarget", async () => {
mockFetchPreviewTarget.mockResolvedValue(makeTarget({ mediaKind: "pdf" }));
mockFetchPreviewFileUrl.mockResolvedValue({
url: "/api/rag/documents/doc-abc/file-signed?token=signed",
expiresAt: 1_700_000_000,
});
await usePreviewStore.getState().open({
documentId: "doc-abc",
backendChunkId: "chunk-xyz",
});
expect(mockFetchPreviewTarget).toHaveBeenCalledWith(
"doc-abc",
"chunk-xyz",
);
});
it("close() revokes the live object URL and clears state", async () => {
mockFetchPreviewTarget.mockResolvedValue(
makeTarget({ mediaKind: "text", filename: "notes.txt" }),
);
mockFetchPreviewFileBlob.mockResolvedValue(makePdfBlob());
await usePreviewStore.getState().open({ documentId: "doc-abc" });
const blobUrl = usePreviewStore.getState().previewBlobUrl;
expect(blobUrl).toMatch(/^blob:/);
usePreviewStore.getState().close();
expect(URL.revokeObjectURL).toHaveBeenCalledWith(blobUrl);
const { previewBlob, previewBlobUrl, target, status } =
usePreviewStore.getState();
expect(previewBlob).toBeNull();
expect(previewBlobUrl).toBeNull();
expect(target).toBeNull();
expect(status).toBe("idle");
});
it("opening doc B revokes doc A's URL before creating doc B's (contracts §5.1)", async () => {
mockFetchPreviewTarget.mockResolvedValue(
makeTarget({ mediaKind: "text", filename: "notes.txt" }),
);
mockFetchPreviewFileBlob.mockResolvedValue(makePdfBlob());
await usePreviewStore.getState().open({ documentId: "doc-A" });
const urlA = usePreviewStore.getState().previewBlobUrl;
expect(urlA).toMatch(/^blob:/);
await usePreviewStore.getState().open({ documentId: "doc-B" });
expect(URL.revokeObjectURL).toHaveBeenCalledWith(urlA);
const urlB = usePreviewStore.getState().previewBlobUrl;
expect(urlB).not.toBe(urlA);
expect(urlB).toMatch(/^blob:/);
});
it("html mediaKind skips blob fetch and sets previewBlobUrl = null (Risk #3)", async () => {
mockFetchPreviewTarget.mockResolvedValue(
makeTarget({ mediaKind: "html", filename: "evil.html" }),
);
await usePreviewStore.getState().open({ documentId: "doc-html" });
expect(mockFetchPreviewFileBlob).not.toHaveBeenCalled();
expect(mockFetchPreviewFileUrl).not.toHaveBeenCalled();
expect(URL.createObjectURL).not.toHaveBeenCalled();
const { previewBlob, previewBlobUrl, status } = usePreviewStore.getState();
expect(previewBlob).toBeNull();
expect(previewBlobUrl).toBeNull();
expect(status).toBe("ready");
});
it("docx mediaKind skips blob fetch and sets previewBlobUrl = null (Risk #3)", async () => {
mockFetchPreviewTarget.mockResolvedValue(
makeTarget({ mediaKind: "docx", filename: "report.docx" }),
);
await usePreviewStore.getState().open({ documentId: "doc-docx" });
expect(mockFetchPreviewFileBlob).not.toHaveBeenCalled();
expect(mockFetchPreviewFileUrl).not.toHaveBeenCalled();
expect(URL.createObjectURL).not.toHaveBeenCalled();
expect(usePreviewStore.getState().previewBlob).toBeNull();
expect(usePreviewStore.getState().previewBlobUrl).toBeNull();
});
it("unknown mediaKind skips blob fetch and sets previewBlobUrl = null", async () => {
mockFetchPreviewTarget.mockResolvedValue(
makeTarget({ mediaKind: "unknown", filename: "data.bin" }),
);
await usePreviewStore.getState().open({ documentId: "doc-bin" });
expect(mockFetchPreviewFileBlob).not.toHaveBeenCalled();
expect(mockFetchPreviewFileUrl).not.toHaveBeenCalled();
expect(URL.createObjectURL).not.toHaveBeenCalled();
expect(usePreviewStore.getState().previewBlob).toBeNull();
expect(usePreviewStore.getState().previewBlobUrl).toBeNull();
});
it("close() when nothing is open does not throw", () => {
expect(() => usePreviewStore.getState().close()).not.toThrow();
});
it("__previewStoreInternals shows no activeBlobUrl after close()", async () => {
mockFetchPreviewTarget.mockResolvedValue(
makeTarget({ mediaKind: "text", filename: "notes.txt" }),
);
mockFetchPreviewFileBlob.mockResolvedValue(makePdfBlob());
await usePreviewStore.getState().open({ documentId: "doc-abc" });
expect(__previewStoreInternals().activeBlobUrl).toMatch(/^blob:/);
usePreviewStore.getState().close();
expect(__previewStoreInternals().activeBlobUrl).toBeNull();
expect(__previewStoreInternals().hasInflightController).toBe(false);
});
it("fetchPreviewTarget error sets status=error and clears target", async () => {
mockFetchPreviewTarget.mockRejectedValue(new Error("404 not found"));
await usePreviewStore.getState().open({ documentId: "missing" });
const { status, error, target } = usePreviewStore.getState();
expect(status).toBe("error");
expect(error).toMatch(/404/);
expect(target).toBeNull();
expect(URL.createObjectURL).not.toHaveBeenCalled();
});
});

View file

@ -0,0 +1,64 @@
import type { PreviewTarget } from "@/features/rag/api/rag-api";
import { PreviewTextView } from "@/features/rag/components/preview-text-view";
import { render, screen } from "@testing-library/react";
import React from "react";
import { describe, expect, it, vi } from "vitest";
vi.mock("@/features/auth", () => ({
authFetch: vi.fn(),
}));
function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
return {
documentId: "doc-abc",
filename: "notes.txt",
contentType: "text/plain",
mediaKind: "text",
byteSize: 100,
status: "completed",
kbId: "kb-1",
threadId: null,
chunkId: "chunk-1",
chunkIndex: 0,
targetPage: 2,
snippet: "alpha\nhighlighted line\nomega",
kind: "text",
imageUrl: null,
sourcePageIndex: null,
pageCharStart: null,
pageCharEnd: null,
lineStart: null,
lineEnd: null,
pdfRegions: [],
...overrides,
};
}
describe("PreviewTextView locator highlight fallback", () => {
it("emphasizes the source excerpt when nullable locators are present", () => {
render(
React.createElement(PreviewTextView, {
target: target({
sourcePageIndex: 1,
pageCharStart: 6,
pageCharEnd: 22,
lineStart: 2,
lineEnd: 2,
}),
}),
);
expect(screen.getByText(/highlighted source excerpt/i)).toBeInTheDocument();
expect(document.querySelector("mark")?.textContent).toContain(
"highlighted line",
);
});
it("keeps the source excerpt visible when locators are missing", () => {
render(React.createElement(PreviewTextView, { target: target() }));
expect(screen.getByText(/source excerpt/i)).toBeInTheDocument();
expect(document.querySelector("mark")).toBeNull();
expect(screen.getByText(/highlighted line/i)).toBeInTheDocument();
});
});

View file

@ -0,0 +1,191 @@
import type { PreviewTarget } from "@/features/rag/api/rag-api";
import { beforeEach, describe, expect, it, vi } from "vitest";
const {
mockAuthFetch,
mockGetAuthToken,
mockEventSourceInstances,
eventSourcePolyfillExport,
authorizationHeader,
} = vi.hoisted(() => ({
mockAuthFetch: vi.fn(),
mockGetAuthToken: vi.fn(),
mockEventSourceInstances: [] as MockEventSource[],
eventSourcePolyfillExport: "EventSourcePolyfill",
authorizationHeader: "Authorization",
}));
vi.mock("@/features/auth", () => ({
authFetch: mockAuthFetch,
getAuthToken: mockGetAuthToken,
}));
interface MockEventSource {
url: string;
options: unknown;
onmessage: ((event: MessageEvent) => void) | null;
onerror: (() => void) | null;
close: ReturnType<typeof vi.fn>;
}
vi.mock("event-source-polyfill", () => ({
[eventSourcePolyfillExport]: class {
url: string;
options: unknown;
onmessage: ((event: MessageEvent) => void) | null = null;
onerror: (() => void) | null = null;
close = vi.fn();
constructor(url: string, options?: unknown) {
this.url = url;
this.options = options;
mockEventSourceInstances.push(this);
}
},
}));
import {
backfillDocumentLocators,
fetchPreviewFileUrl,
fetchPreviewTarget,
subscribeToJobEvents,
} from "@/features/rag/api/rag-api";
function target(): PreviewTarget {
return {
documentId: "doc-abc",
filename: "report.pdf",
contentType: "application/pdf",
mediaKind: "pdf",
byteSize: 100,
status: "completed",
kbId: "kb-1",
threadId: null,
chunkId: "chunk-xyz",
chunkIndex: 0,
targetPage: 1,
snippet: "excerpt",
kind: "text",
imageUrl: null,
sourcePageIndex: 0,
pageCharStart: 0,
pageCharEnd: 7,
lineStart: 1,
lineEnd: 1,
pdfRegions: [],
};
}
beforeEach(() => {
mockAuthFetch.mockReset();
mockGetAuthToken.mockReset();
mockEventSourceInstances.length = 0;
});
describe("RAG API preview target", () => {
it("URL-encodes documentId and chunk_id", async () => {
mockAuthFetch.mockResolvedValue(
new Response(JSON.stringify(target()), {
status: 200,
headers: { "Content-Type": "application/json" },
}),
);
await fetchPreviewTarget("doc id/with?slash", "chunk id/with?amp&eq=1");
expect(mockAuthFetch).toHaveBeenCalledWith(
"/api/rag/documents/doc%20id%2Fwith%3Fslash/preview-target?chunk_id=chunk%20id%2Fwith%3Famp%26eq%3D1",
);
});
it("fetches signed preview URL without adding a bearer token query", async () => {
mockAuthFetch.mockResolvedValue(
new Response(
JSON.stringify({
url: "/api/rag/documents/doc-abc/file-signed?token=signed-preview",
expiresAt: 1_700_000_000,
}),
{
status: 200,
headers: { "Content-Type": "application/json" },
},
),
);
const result = await fetchPreviewFileUrl("doc id/with?slash");
expect(mockAuthFetch).toHaveBeenCalledWith(
"/api/rag/documents/doc%20id%2Fwith%3Fslash/file-url",
undefined,
);
expect(result.url).toContain("token=signed-preview");
expect(result.url).not.toContain("Bearer");
expect(result.url).not.toContain("Authorization");
});
it("posts the explicit locator backfill action", async () => {
mockAuthFetch.mockResolvedValue(
new Response(
JSON.stringify({
documentId: "doc-abc",
totalChunks: 1,
matched: 1,
alreadyLocated: 0,
ambiguous: 0,
missing: 0,
skipped: 0,
regionsMatched: 0,
pagesRefreshed: 1,
}),
{
status: 200,
headers: { "Content-Type": "application/json" },
},
),
);
await backfillDocumentLocators("doc id/with?slash");
expect(mockAuthFetch).toHaveBeenCalledWith(
"/api/rag/documents/doc%20id%2Fwith%3Fslash/locators/backfill",
{ method: "POST" },
);
});
});
describe("RAG API job events", () => {
it("opens SSE with Authorization header instead of token query params", () => {
mockGetAuthToken.mockReturnValue("mock-token-123");
const unsubscribe = subscribeToJobEvents("job id/with?slash", {});
expect(mockEventSourceInstances).toHaveLength(1);
const source = mockEventSourceInstances[0];
expect(source.url).toContain(
"/api/rag/jobs/job%20id%2Fwith%3Fslash/events",
);
expect(source.url).not.toContain("token=");
expect(source.options).toEqual({
headers: {
[authorizationHeader]: "Bearer mock-token-123",
},
});
unsubscribe();
expect(source.close).toHaveBeenCalled();
});
it("omits EventSource options when there is no bearer token", () => {
mockGetAuthToken.mockReturnValue(null);
const unsubscribe = subscribeToJobEvents("job-abc", {});
expect(mockEventSourceInstances).toHaveLength(1);
const source = mockEventSourceInstances[0];
expect(source.url).toContain("/api/rag/jobs/job-abc/events");
expect(source.url).not.toContain("token=");
expect(source.options).toBeUndefined();
unsubscribe();
});
});

View file

@ -0,0 +1,84 @@
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import React from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
const { mockOpenPreview } = vi.hoisted(() => ({
mockOpenPreview: vi.fn(),
}));
vi.mock("@/features/rag/stores/preview-store", () => ({
usePreviewStore: (
selector?: (state: { open: typeof mockOpenPreview }) => unknown,
) => {
const state = { open: mockOpenPreview };
return typeof selector === "function" ? selector(state) : state;
},
}));
vi.mock("@assistant-ui/react", () => ({
useAuiState: (
selector: (state: { message: { content: unknown[] } }) => unknown,
) =>
selector({
message: { content: [{ type: "text", text: "Answer ready." }] },
}),
}));
import { SearchKnowledgeBaseToolUI } from "@/components/assistant-ui/tool-ui-search-knowledge-base";
const TOOL_UI = SearchKnowledgeBaseToolUI as React.ComponentType<
Record<string, unknown>
>;
const SEARCHED_DOCS_BUTTON_RE = /searched docs/i;
const MAIN_PREVIEW_BUTTON_RE = /open preview of main\.pdf/i;
const LEGACY_PREVIEW_BUTTON_RE = /open preview of legacy\.pdf/i;
function renderTool(result: string) {
return render(
React.createElement(TOOL_UI, {
args: { query: "what is interior modeling?" },
result,
status: { type: "complete" },
}),
);
}
beforeEach(() => {
mockOpenPreview.mockClear();
});
describe("SearchKnowledgeBaseToolUI preview routing", () => {
it("opens preview from a retrieved chunk source label", async () => {
renderTool(
'<chunk id="1" source="main.pdf" page="3" document_id="doc-abc" chunk_id="chunk-xyz">The paper objectives.</chunk>',
);
await userEvent.click(
screen.getByRole("button", { name: SEARCHED_DOCS_BUTTON_RE }),
);
await userEvent.click(
screen.getByRole("button", { name: MAIN_PREVIEW_BUTTON_RE }),
);
expect(mockOpenPreview).toHaveBeenCalledWith({
documentId: "doc-abc",
backendChunkId: "chunk-xyz",
});
});
it("keeps legacy chunk labels non-clickable without durable IDs", async () => {
renderTool(
'<chunk id="1" source="legacy.pdf" page="3">Legacy chunk text.</chunk>',
);
await userEvent.click(
screen.getByRole("button", { name: SEARCHED_DOCS_BUTTON_RE }),
);
expect(
screen.queryByRole("button", { name: LEGACY_PREVIEW_BUTTON_RE }),
).toBeNull();
expect(mockOpenPreview).not.toHaveBeenCalled();
});
});

View file

@ -0,0 +1,80 @@
import { fireEvent, render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import React from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
const { mockOpen } = vi.hoisted(() => ({ mockOpen: vi.fn() }));
vi.mock("@/features/rag/stores/preview-store", () => ({
usePreviewStore: (
selector?: (state: { open: typeof mockOpen }) => unknown,
) => {
const state = { open: mockOpen };
return typeof selector === "function" ? selector(state) : state;
},
}));
import { DocumentSourceBadge } from "@/components/assistant-ui/sources";
function source(overrides: Record<string, unknown> = {}) {
return {
kind: "document" as const,
chunkId: "2",
documentId: "doc-abc",
backendChunkId: "chunk-xyz",
filename: "report.pdf",
page: "7",
score: "0.85",
text: "The margin rose to 18%.",
...overrides,
};
}
beforeEach(() => {
mockOpen.mockClear();
});
describe("DocumentSourceBadge preview routing", () => {
it("opens preview with durable document and backend chunk IDs on click", async () => {
render(React.createElement(DocumentSourceBadge, { source: source() }));
await userEvent.click(
screen.getByRole("button", { name: /open preview/i }),
);
expect(mockOpen).toHaveBeenCalledWith({
documentId: "doc-abc",
backendChunkId: "chunk-xyz",
});
});
it("opens preview from Enter and Space", () => {
render(React.createElement(DocumentSourceBadge, { source: source() }));
const badge = screen.getByRole("button", { name: /open preview/i });
fireEvent.keyDown(badge, { key: "Enter" });
fireEvent.keyDown(badge, { key: " " });
expect(mockOpen).toHaveBeenCalledTimes(2);
});
it("legacy source without durable IDs remains hover-only", async () => {
render(
React.createElement(DocumentSourceBadge, {
source: source({ documentId: null, backendChunkId: null }),
}),
);
expect(screen.queryByRole("button", { name: /open preview/i })).toBeNull();
await userEvent.click(screen.getByText("[2]"));
expect(mockOpen).not.toHaveBeenCalled();
});
it("applies brand-aligned interactive styling when preview is clickable", () => {
render(React.createElement(DocumentSourceBadge, { source: source() }));
const badge = screen.getByRole("button", { name: /open preview/i });
expect(badge).toHaveClass("cursor-pointer");
expect(badge).toHaveClass("hover:bg-chat-icon-bg-hover!");
expect(badge).toHaveClass("focus-visible:ring-ring/50");
});
});

View file

@ -1,24 +1,26 @@
"use client";
import { openLink } from "@/lib/open-link";
import {
memo,
useState,
useRef,
useEffect,
useCallback,
type ComponentProps,
type FC,
} from "react";
import { FileTextIcon } from "lucide-react";
import { useMessage } from "@assistant-ui/react";
import { cn } from "@/lib/utils";
import { Badge, badgeVariants, type BadgeProps } from "./badge";
import {
HoverCard,
HoverCardTrigger,
HoverCardContent,
HoverCardTrigger,
} from "@/components/ui/hover-card";
import { usePreviewStore } from "@/features/rag/stores/preview-store";
import { openLink } from "@/lib/open-link";
import { cn } from "@/lib/utils";
import { useMessage } from "@assistant-ui/react";
import { FileTextIcon } from "lucide-react";
import {
type ComponentProps,
type FC,
type KeyboardEvent as ReactKeyboardEvent,
memo,
useCallback,
useEffect,
useRef,
useState,
} from "react";
import { Badge, type BadgeProps, badgeVariants } from "./badge";
// ── Helpers ──────────────────────────────────────────────────
@ -45,8 +47,12 @@ function SourceIcon({
}: ComponentProps<"span"> & { url: string; size?: number }) {
const [hasError, setHasError] = useState(false);
const domain = extractDomain(url);
const SIZE_CLASSES: Record<number, string> = { 3: "size-3", 4: "size-4", 5: "size-5" };
const sizeClass = SIZE_CLASSES[size] ?? "size-3";
const sizeClasses: Record<number, string> = {
3: "size-3",
4: "size-4",
5: "size-5",
};
const sizeClass = sizeClasses[size] ?? "size-3";
if (hasError) {
return (
@ -101,7 +107,7 @@ function Source({
}: SourceProps) {
return (
<Badge
asChild
asChild={true}
variant={variant}
size={size}
className={cn(
@ -136,7 +142,12 @@ interface UrlSourceData {
interface DocSourceData {
kind: "document";
/** Visible citation id (the `[N]` reference). Display only. */
chunkId: string;
/** Durable backend `rag_documents.id`. null on legacy sources. */
documentId: string | null;
/** Durable backend `rag_chunks.id`. null on legacy sources. */
backendChunkId: string | null;
filename: string;
page?: string;
text: string;
@ -145,9 +156,7 @@ interface DocSourceData {
type SourceData = UrlSourceData | DocSourceData;
function sourceKey(source: SourceData): string {
return source.kind === "url"
? `url:${source.url}`
: `doc:${source.chunkId}`;
return source.kind === "url" ? `url:${source.url}` : `doc:${source.chunkId}`;
}
const SourceBadge: FC<{ source: UrlSourceData }> = ({ source }) => {
@ -156,7 +165,7 @@ const SourceBadge: FC<{ source: UrlSourceData }> = ({ source }) => {
return (
<HoverCard openDelay={0} closeDelay={0}>
<HoverCardTrigger asChild>
<HoverCardTrigger asChild={true}>
<span className="inline-block">
<Source href={source.url}>
<SourceIcon url={source.url} />
@ -193,13 +202,53 @@ const DocumentSourceBadge: FC<{ source: DocSourceData }> = ({ source }) => {
const metaParts: string[] = [];
if (source.page) metaParts.push(`page ${source.page}`);
// Preview is clickable IFF both durable IDs are present (contracts
// §4.1 routing rule + Q3). Legacy sources fall through to a
// non-interactive badge with hover-only behavior.
const isClickable =
source.documentId !== null && source.backendChunkId !== null;
const openPreview = usePreviewStore((s) => s.open);
const handleOpen = useCallback(() => {
if (!isClickable || !source.documentId) return;
void openPreview({
documentId: source.documentId,
backendChunkId: source.backendChunkId,
});
}, [isClickable, openPreview, source.documentId, source.backendChunkId]);
const handleKeyDown = useCallback(
(e: ReactKeyboardEvent<HTMLElement>) => {
if (!isClickable) return;
if (e.key === "Enter" || e.key === " ") {
e.preventDefault();
handleOpen();
}
},
[isClickable, handleOpen],
);
return (
<HoverCard openDelay={0} closeDelay={0}>
<HoverCardTrigger asChild>
<HoverCardTrigger asChild={true}>
<span className="inline-block">
<Badge
variant="outline"
className="rounded-full cursor-default inline-flex items-center gap-1.5 outline-none"
{...(isClickable
? {
role: "button",
tabIndex: 0,
onClick: handleOpen,
onKeyDown: handleKeyDown,
"aria-label": `Open preview of ${source.filename}`,
}
: {})}
className={cn(
"rounded-full inline-flex items-center gap-1.5 outline-none",
isClickable
? "cursor-pointer hover:bg-chat-icon-bg-hover! hover:text-chat-icon-fg-hover! focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50"
: "cursor-default",
)}
>
<span className="font-mono text-[10px] font-semibold text-muted-foreground">
[{source.chunkId}]
@ -229,6 +278,11 @@ const DocumentSourceBadge: FC<{ source: DocSourceData }> = ({ source }) => {
<p className="text-xs text-white/70 leading-relaxed line-clamp-3 whitespace-pre-wrap">
{source.text}
</p>
{isClickable ? (
<p className="mt-1 text-[10px] text-white/50">
Click to open preview
</p>
) : null}
</div>
</div>
</HoverCardContent>
@ -269,6 +323,8 @@ const SourcesGroup: FC = () => {
) {
const docPart = part as {
chunkId?: string;
documentId?: string | null;
backendChunkId?: string | null;
filename?: string;
page?: string;
text?: string;
@ -277,6 +333,8 @@ const SourcesGroup: FC = () => {
sources.push({
kind: "document",
chunkId: docPart.chunkId,
documentId: docPart.documentId ?? null,
backendChunkId: docPart.backendChunkId ?? null,
filename: docPart.filename,
page: docPart.page,
text: docPart.text ?? "",
@ -340,7 +398,7 @@ const SourcesGroup: FC = () => {
{/* Hidden measurement container — renders all badges to measure row positions */}
<div
ref={containerRef}
aria-hidden
aria-hidden={true}
className="flex w-full flex-wrap gap-1 invisible absolute pointer-events-none"
>
{sources.map((source) => (

View file

@ -3,14 +3,15 @@
"use client";
import { authFetch } from "@/features/auth";
import { usePreviewStore } from "@/features/rag/stores/preview-store";
import { cn } from "@/lib/utils";
import {
type ToolCallMessagePartComponent,
useAuiState,
} from "@assistant-ui/react";
import { authFetch } from "@/features/auth";
import { FileTextIcon, ImageIcon, LoaderIcon } from "lucide-react";
import { memo, useEffect, useState } from "react";
import { cn } from "@/lib/utils";
import { memo, useCallback, useEffect, useState } from "react";
import {
ToolFallbackContent,
ToolFallbackRoot,
@ -18,14 +19,29 @@ import {
} from "./tool-fallback";
export interface ParsedChunk {
/** Visible citation id the model uses inside `[N]` references. Display
* only; never sent to the backend as a chunk_id. */
id: string;
source: string;
page?: string;
chunkIndex?: string;
tokens?: string;
sourcePageIndex?: string;
pageCharStart?: string;
pageCharEnd?: string;
lineStart?: string;
lineEnd?: string;
kind?: string;
imageUrl?: string;
text: string;
/** Durable `rag_documents.id`. Carries through when the tool XML
* includes `document_id="..."`. Absent on legacy tool output. */
documentId?: string;
/** Durable `rag_chunks.id`. Carries through when the tool XML
* includes `chunk_id="..."`. Absent on legacy tool output. The
* preview routing value sent as `?chunk_id=` to `/preview-target`;
* never the same as the visible `id`. */
backendChunkId?: string;
}
const ATTR_RE = /(\w+)="([^"]*)"/g;
@ -59,9 +75,17 @@ export function parseChunks(raw: string): ParsedChunk[] {
page: attrs.page,
chunkIndex: attrs.chunk_index,
tokens: attrs.tokens,
sourcePageIndex: attrs.source_page_index,
pageCharStart: attrs.page_char_start,
pageCharEnd: attrs.page_char_end,
lineStart: attrs.line_start,
lineEnd: attrs.line_end,
kind: attrs.kind,
imageUrl: attrs.image_url,
text,
// Durable backend ids (legacy XML omits both → preview gated off).
...(attrs.document_id ? { documentId: attrs.document_id } : {}),
...(attrs.chunk_id ? { backendChunkId: attrs.chunk_id } : {}),
});
}
match = CHUNK_RE.exec(raw);
@ -82,9 +106,11 @@ function useAuthedImageUrl(path: string | undefined): string | undefined {
}
let cancelled = false;
let objectUrl: string | undefined;
void authFetch(path)
.then(async (response) => {
if (!response.ok) throw new Error(`image fetch ${response.status}`);
authFetch(path)
.then((response) => {
if (!response.ok) {
throw new Error(`image fetch ${response.status}`);
}
return response.blob();
})
.then((blob) => {
@ -123,31 +149,66 @@ function ChunkImage({ url, alt }: { url: string; alt: string }) {
}
function ChunkCard({ chunk }: { chunk: ParsedChunk }) {
const openPreview = usePreviewStore((s) => s.open);
const meta: string[] = [];
if (chunk.page) meta.push(`page ${chunk.page}`);
if (chunk.tokens) meta.push(`${chunk.tokens} tok`);
if (chunk.chunkIndex) meta.push(`#${chunk.chunkIndex}`);
if (chunk.kind && chunk.kind !== "text") meta.push(chunk.kind);
const documentId = chunk.documentId;
const backendChunkId = chunk.backendChunkId;
const isPreviewable = Boolean(documentId && backendChunkId);
const handleOpenPreview = useCallback(() => {
if (!(documentId && backendChunkId)) {
return;
}
Promise.resolve(
openPreview({
documentId,
backendChunkId,
}),
).catch(() => undefined);
}, [backendChunkId, documentId, openPreview]);
const sourceLabel = (
<>
<span className="rounded bg-foreground/10 px-1.5 py-0.5 font-mono text-[10px] font-semibold">
[{chunk.id}]
</span>
{chunk.kind === "image" ? (
<ImageIcon className="size-3 shrink-0 text-muted-foreground" />
) : (
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
)}
<span className="truncate font-medium" title={chunk.source}>
{chunk.source}
</span>
</>
);
return (
<div
data-slot="rag-chunk-card"
className="rounded-md border border-foreground/10 bg-muted/40 p-2.5 text-xs"
>
<div className="mb-1.5 flex items-center justify-between gap-2">
<div className="flex min-w-0 items-center gap-1.5">
<span className="rounded bg-foreground/10 px-1.5 py-0.5 font-mono text-[10px] font-semibold">
[{chunk.id}]
</span>
{chunk.kind === "image" ? (
<ImageIcon className="size-3 shrink-0 text-muted-foreground" />
) : (
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
)}
<span className="truncate font-medium" title={chunk.source}>
{chunk.source}
</span>
</div>
{isPreviewable ? (
<button
type="button"
className={cn(
"flex min-w-0 cursor-pointer items-center gap-1.5 rounded-sm text-left outline-none transition-colors",
"hover:text-primary focus-visible:ring-[3px] focus-visible:ring-ring/50",
)}
onClick={handleOpenPreview}
aria-label={`Open preview of ${chunk.source}`}
title="Open preview"
>
{sourceLabel}
</button>
) : (
<div className="flex min-w-0 items-center gap-1.5">{sourceLabel}</div>
)}
{meta.length > 0 ? (
<span className="shrink-0 text-[10px] tabular-nums text-muted-foreground">
{meta.join(" · ")}

View file

@ -1,7 +1,18 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
type ParsedChunk,
parseChunks,
} from "@/components/assistant-ui/tool-ui-search-knowledge-base";
import { getAuthToken } from "@/features/auth/session";
import {
type SearchHit,
type SearchRequest,
listKBDocuments,
listThreadDocuments,
search as ragSearch,
} from "@/features/rag/api/rag-api";
import { apiUrl } from "@/lib/api-base";
import { toast } from "@/lib/toast";
import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core";
@ -37,12 +48,12 @@ import type {
OpenAIMessageContent,
} from "../types/api";
import type { ChatModelSummary } from "../types/runtime";
import { getImageInputUnavailableReason } from "../utils/image-input-support";
import {
getStoredChatThread,
listStoredChatThreads,
updateStoredChatThread,
} from "../utils/chat-history-storage";
import { getImageInputUnavailableReason } from "../utils/image-input-support";
import {
hasClosedThinkTag,
parseAssistantContent,
@ -56,17 +67,6 @@ import {
streamChatCompletions,
validateModel,
} from "./chat-api";
import {
type SearchHit,
type SearchRequest,
listKBDocuments,
listThreadDocuments,
search as ragSearch,
} from "@/features/rag/api/rag-api";
import {
type ParsedChunk,
parseChunks,
} from "@/components/assistant-ui/tool-ui-search-knowledge-base";
import type { RagMode, RagSource } from "./chat-settings-api";
import {
createOpenAIContainer,
@ -123,8 +123,7 @@ function buildRagRequest(
function formatRagContext(hits: SearchHit[]): string {
const parts = hits.map((h) => {
const name = h.filename ?? `chunk ${h.chunk_index}`;
const pageAttr =
h.page_number != null ? ` page="${h.page_number}"` : "";
const pageAttr = h.page_number != null ? ` page="${h.page_number}"` : "";
return `<source filename="${name}"${pageAttr}>\n${h.text}\n</source>`;
});
return `<context>\nThe following documents may help answer the user's question:\n${parts.join("\n")}\n</context>`;
@ -251,9 +250,27 @@ interface DocumentSourcePart {
type: "source";
sourceType: "document";
id: string;
/** Display alias of `citationId`. Kept so the existing sources.tsx
* renderer keeps working; new code SHOULD use `citationId`. NEVER
* sent to backend as the durable chunk_id. */
chunkId: string;
/** Visible model-citation id (the `[N]` reference). Display only. */
citationId: string;
/** Durable `rag_documents.id` from tool XML `document_id=`. Null
* when the source came from legacy XML lacking the attribute;
* preview routing is gated off in that case. */
documentId: string | null;
/** Durable `rag_chunks.id` from tool XML `chunk_id=`. Null on
* legacy XML. Sent as `?chunk_id=` to `/preview-target`. */
backendChunkId: string | null;
filename: string;
page?: string;
sourcePageIndex?: string;
pageCharStart?: string;
pageCharEnd?: string;
lineStart?: string;
lineEnd?: string;
score?: string;
text: string;
}
@ -273,32 +290,93 @@ function extractCitedIds(text: string): Set<string> {
return ids;
}
/** Build doc-shaped source parts for chunks the model actually cited.
function indexChunksByCitationId(
allChunks: ParsedChunk[],
): Map<string, ParsedChunk> {
const byId = new Map<string, ParsedChunk>();
for (const chunk of allChunks) {
if (!byId.has(chunk.id)) {
byId.set(chunk.id, chunk);
}
}
return byId;
}
function documentSourceIds(
allChunks: ParsedChunk[],
citedIds: Set<string>,
): string[] {
if (citedIds.size > 0) {
return Array.from(citedIds);
}
return allChunks.map((chunk) => chunk.id);
}
function toDocumentSourcePart(
id: string,
chunk: ParsedChunk,
): DocumentSourcePart {
const part: DocumentSourcePart = {
type: "source",
sourceType: "document",
id: `rag-${id}`,
chunkId: id,
citationId: id,
documentId: chunk.documentId ?? null,
backendChunkId: chunk.backendChunkId ?? null,
filename: chunk.source,
text: chunk.text,
};
if (chunk.page) {
part.page = chunk.page;
}
if (chunk.sourcePageIndex) {
part.sourcePageIndex = chunk.sourcePageIndex;
}
if (chunk.pageCharStart) {
part.pageCharStart = chunk.pageCharStart;
}
if (chunk.pageCharEnd) {
part.pageCharEnd = chunk.pageCharEnd;
}
if (chunk.lineStart) {
part.lineStart = chunk.lineStart;
}
if (chunk.lineEnd) {
part.lineEnd = chunk.lineEnd;
}
if (chunk.score) {
part.score = chunk.score;
}
return part;
}
/** Build doc-shaped source parts for chunks the model cited.
* `allChunks` is the flat union of every search_knowledge_base tool
* result in this turn (deduped by id). Returns one part per unique
* cited id that maps to a real chunk; hallucinated `[99]` refs without
* a matching chunk are silently dropped. */
* result in this turn (deduped by id). If the model forgets literal
* `[N]` ids, fall back to retrieved chunks so source chips remain
* visible and previewable. Hallucinated `[99]` refs without a matching
* chunk are silently dropped. */
function buildDocumentSourceParts(
allChunks: ParsedChunk[],
citedIds: Set<string>,
): DocumentSourcePart[] {
const byId = new Map<string, ParsedChunk>();
for (const chunk of allChunks) {
if (!byId.has(chunk.id)) byId.set(chunk.id, chunk);
}
const byId = indexChunksByCitationId(allChunks);
const idsToShow = documentSourceIds(allChunks, citedIds);
const out: DocumentSourcePart[] = [];
for (const id of citedIds) {
const emittedIds = new Set<string>();
for (const id of idsToShow) {
if (emittedIds.has(id)) {
continue;
}
const chunk = byId.get(id);
if (!chunk) continue;
out.push({
type: "source",
sourceType: "document",
id: `rag-${id}`,
chunkId: id,
filename: chunk.source,
...(chunk.page ? { page: chunk.page } : {}),
text: chunk.text,
});
if (!chunk) {
continue;
}
emittedIds.add(id);
out.push(toDocumentSourcePart(id, chunk));
}
return out;
}
@ -954,7 +1032,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// Re-read store after potential auto-load / model ready wait
runtime = useChatRuntimeStore.getState();
const { params } = runtime;
const { supportsTools, toolsEnabled, codeToolsEnabled, imageToolsEnabled } = runtime;
const {
supportsTools,
toolsEnabled,
codeToolsEnabled,
imageToolsEnabled,
} = runtime;
const externalSelection = parseExternalModelId(params.checkpoint);
const isExternalRequest = externalSelection !== null;
if (
@ -993,33 +1076,30 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
throw new Error("Missing connection API key.");
}
const webSearchEnabledForThisTurn =
Boolean(
externalProvider &&
toolsEnabled &&
providerSupportsBuiltinWebSearch(externalProvider.providerType),
);
const codeExecEnabledForThisTurn =
Boolean(
externalProvider &&
externalSelection &&
codeToolsEnabled &&
providerSupportsBuiltinCodeExecution(
externalProvider.providerType,
externalSelection.modelId,
externalProvider.baseUrl,
),
);
const webSearchEnabledForThisTurn = Boolean(
externalProvider &&
toolsEnabled &&
providerSupportsBuiltinWebSearch(externalProvider.providerType),
);
const codeExecEnabledForThisTurn = Boolean(
externalProvider &&
externalSelection &&
codeToolsEnabled &&
providerSupportsBuiltinCodeExecution(
externalProvider.providerType,
externalSelection.modelId,
externalProvider.baseUrl,
),
);
// web_fetch shares the Search pill with web_search (no separate
// UI toggle), so it follows toolsEnabled. Anthropic is the only
// provider that ships it today; on others providerSupportsBuiltinWebFetch
// returns false and this stays inert.
const webFetchEnabledForThisTurn =
Boolean(
externalProvider &&
toolsEnabled &&
providerSupportsBuiltinWebFetch(externalProvider.providerType),
);
const webFetchEnabledForThisTurn = Boolean(
externalProvider &&
toolsEnabled &&
providerSupportsBuiltinWebFetch(externalProvider.providerType),
);
const providerShipsWebFetch = Boolean(
externalProvider &&
providerSupportsBuiltinWebFetch(externalProvider.providerType),
@ -1049,7 +1129,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// entirely so retrieval only happens via the LLM-invoked
// search_knowledge_base tool. Flip back to true to restore the
// always-on grounding for external providers / non-tool models.
const RAG_PREFETCH_ENABLED = false;
const ragPrefetchEnabled = false;
const ragSource = runtime.ragSource;
const ragToolEnabled = runtime.ragToolEnabled;
@ -1115,13 +1195,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
});
}
if (RAG_PREFETCH_ENABLED && ragToolEnabled && ragSource.kind !== "off") {
if (ragPrefetchEnabled && ragToolEnabled && ragSource.kind !== "off") {
const lastUser = [...outboundMessages]
.reverse()
.find((m) => m.role === "user");
const queryText = lastUser
? extractMessageText(lastUser.content)
: "";
const queryText = lastUser ? extractMessageText(lastUser.content) : "";
if (queryText.trim()) {
const ragReq = buildRagRequest(
ragSource,
@ -1145,8 +1223,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
outboundMessages[0]?.role === "system" &&
typeof outboundMessages[0].content === "string"
) {
outboundMessages[0].content =
`${block}\n\n${outboundMessages[0].content}`;
outboundMessages[0].content = `${block}\n\n${outboundMessages[0].content}`;
} else {
outboundMessages.unshift({ role: "system", content: block });
}
@ -1454,7 +1531,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
? reasoningEffort
: "low";
const externalReasoningEnabled =
!externalReasoningCaps.supportsReasoningOff ? true : reasoningEnabled;
externalReasoningCaps.supportsReasoningOff ? reasoningEnabled : true;
const buildRequestPayload = async (
forceRefreshPublicKey = false,
): Promise<OpenAIChatCompletionsRequest> => {
@ -1540,8 +1617,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
) {
void updateStoredChatThreadEventually(t.id, {
openaiCodeExecContainerId: null,
})
.catch(() => {});
}).catch(() => {});
continue;
}
openaiCodeExecContainerId = t.openaiCodeExecContainerId;
@ -1585,8 +1661,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
openaiCodeExecContainerId = created.id;
void updateStoredChatThreadEventually(resolvedThreadId, {
openaiCodeExecContainerId: created.id,
})
.catch(() => {});
}).catch(() => {});
} catch {
// Fall back to backend's container_auto path on
// failure — keeps the chat moving; the next turn
@ -1699,7 +1774,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// attaches `cache_control.ttl` when the value is one of
// "5m" / "1h" (see external_provider.py near line 1375),
// so unknown values are a no-op end-to-end.
...(supportsProviderPromptCacheTtl(externalProvider.providerType) &&
...(supportsProviderPromptCacheTtl(
externalProvider.providerType,
) &&
(externalProvider.enablePromptCaching ?? true) &&
isPromptCacheTtl(externalProvider.promptCacheTtl)
? { prompt_cache_ttl: externalProvider.promptCacheTtl }
@ -1744,8 +1821,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
...(supportsPreserveThinking
? { preserve_thinking: preserveThinking }
: {}),
...(supportsTools
&& (toolsEnabled || codeToolsEnabled || ragToolPathTaken)
...(supportsTools &&
(toolsEnabled || codeToolsEnabled || ragToolPathTaken)
? {
enable_tools: true,
enabled_tools: [
@ -1760,9 +1837,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
? {
rag_scope: {
kb_id:
ragSource.kind === "kb"
? ragSource.kbId
: null,
ragSource.kind === "kb" ? ragSource.kbId : null,
thread_id:
ragSource.kind === "thread"
? (resolvedThreadId ?? null)
@ -1840,8 +1915,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
: "openaiCodeExecContainerId";
void updateStoredChatThreadEventually(resolvedThreadId, {
[field]: null,
})
.catch(() => {});
}).catch(() => {});
}
continue;
}
@ -2026,11 +2100,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
}
if (reasoning) {
if (!reasoningContentOpen) {
if (reasoningContentOpen) {
cumulativeText += reasoning;
} else {
cumulativeText += `<think>${reasoning}`;
reasoningContentOpen = true;
} else {
cumulativeText += reasoning;
}
}
if (delta) {
@ -2125,8 +2199,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
});
// RAG: flatten chunks across every search_knowledge_base call this
// turn, then emit doc-source parts only for ids the model actually
// cited as [N] in its final reply.
// turn, then emit previewable doc-source chips for cited chunks.
// If the model omits literal [N] ids, show the retrieved chunks so
// the answer still has a visible citation/preview affordance.
const ragChunks = toolCallParts.flatMap((tc) => {
if (tc.toolName !== "search_knowledge_base" || !tc.result) {
return [];

File diff suppressed because it is too large Load diff

View file

@ -4,6 +4,7 @@
import { authFetch, getAuthToken } from "@/features/auth";
import { apiUrl } from "@/lib/api-base";
import { formatFastApiDetail } from "@/lib/format-fastapi-error";
import { EventSourcePolyfill } from "event-source-polyfill";
export type ChunkingStrategy = "standard" | "late";
export type KBMode = "text" | "multimodal";
@ -47,6 +48,74 @@ export interface SearchHit {
filename: string | null;
kind?: "text" | "image" | "caption";
image_url?: string | null;
source_page_index?: number | null;
page_char_start?: number | null;
page_char_end?: number | null;
line_start?: number | null;
line_end?: number | null;
}
// --- Preview target (durable backend-id routing) ---
export type PreviewMediaKind =
| "pdf"
| "text"
| "docx"
| "html"
| "image"
| "unknown";
export type PreviewChunkKind = "text" | "image" | "caption";
export interface PreviewPdfRegion {
pageIndex: number;
pageNumber: number | null;
x: number;
y: number;
width: number;
height: number;
confidence: "exact";
source: string;
}
export interface PreviewTarget {
documentId: string;
filename: string;
contentType: string | null;
mediaKind: PreviewMediaKind;
byteSize: number;
status: string;
kbId: string | null;
threadId: string | null;
chunkId: string | null;
chunkIndex: number | null;
targetPage: number | null;
snippet: string | null;
kind: PreviewChunkKind | null;
imageUrl: string | null;
sourcePageIndex: number | null;
pageCharStart: number | null;
pageCharEnd: number | null;
lineStart: number | null;
lineEnd: number | null;
pdfRegions: PreviewPdfRegion[];
}
export interface PreviewFileUrl {
url: string;
expiresAt: number;
}
export interface LocatorBackfillResult {
documentId: string;
totalChunks: number;
matched: number;
alreadyLocated: number;
ambiguous: number;
missing: number;
skipped: number;
regionsMatched: number;
pagesRefreshed: number;
}
export interface SearchRequest {
@ -63,7 +132,13 @@ export interface SearchRequest {
}
export type JobEvent =
| { type: "status"; status: string; stage?: string | null; progress?: number; error?: string | null }
| {
type: "status";
status: string;
stage?: string | null;
progress?: number;
error?: string | null;
}
| { type: "progress"; stage: string; progress: number }
| { type: "complete"; num_chunks: number }
| { type: "error"; error: string };
@ -72,7 +147,9 @@ function parseErrorText(status: number, body: unknown): string {
if (body && typeof body === "object") {
const detail = (body as { detail?: unknown }).detail;
const formatted = formatFastApiDetail(detail);
if (formatted) return formatted;
if (formatted) {
return formatted;
}
}
return `Request failed (${status})`;
}
@ -95,7 +172,9 @@ async function throwOnError(response: Response): Promise<void> {
export async function listKnowledgeBases(): Promise<KnowledgeBase[]> {
const response = await authFetch("/api/rag/knowledge-bases");
const body = await parseJsonOrThrow<{ knowledge_bases: KnowledgeBase[] }>(response);
const body = await parseJsonOrThrow<{ knowledge_bases: KnowledgeBase[] }>(
response,
);
return body.knowledge_bases;
}
@ -136,7 +215,9 @@ export async function listKBDocuments(kbId: string): Promise<RagDocument[]> {
return body.documents;
}
export async function listThreadDocuments(threadId: string): Promise<RagDocument[]> {
export async function listThreadDocuments(
threadId: string,
): Promise<RagDocument[]> {
const response = await authFetch(
`/api/rag/threads/${encodeURIComponent(threadId)}/documents`,
);
@ -187,7 +268,9 @@ export interface ThreadIndexSummary {
export async function listThreadIndexes(): Promise<ThreadIndexSummary[]> {
const response = await authFetch("/api/rag/thread-indexes");
const body = await parseJsonOrThrow<{ threads: ThreadIndexSummary[] }>(response);
const body = await parseJsonOrThrow<{ threads: ThreadIndexSummary[] }>(
response,
);
return body.threads;
}
@ -342,7 +425,8 @@ export async function search(req: SearchRequest): Promise<SearchHit[]> {
// --- Ingestion SSE ---
/** Subscribe to a job's SSE stream; returns an unsubscribe fn.
* Token goes via ?token= since EventSource cannot send headers. */
* Use the EventSource polyfill so the bearer token rides in an
* Authorization header instead of leaking through URL query params. */
export function subscribeToJobEvents(
jobId: string,
handlers: {
@ -352,9 +436,17 @@ export function subscribeToJobEvents(
},
): () => void {
const token = getAuthToken();
const params = token ? `?token=${encodeURIComponent(token)}` : "";
const url = apiUrl(`/api/rag/jobs/${encodeURIComponent(jobId)}/events${params}`);
const source = new EventSource(url);
const url = apiUrl(`/api/rag/jobs/${encodeURIComponent(jobId)}/events`);
const source = new EventSourcePolyfill(
url,
token
? {
headers: {
Authorization: `Bearer ${token}`,
},
}
: undefined,
);
source.onmessage = (e) => {
try {
@ -381,3 +473,68 @@ export function subscribeToJobEvents(
handlers.onClose?.();
};
}
// --- Preview target + blob ---
/** Fetch the preview-target metadata for a document. When `chunkId` is
* provided it must belong to `documentId`; mismatch collapses to 404
* with the same "Document not found" body as missing/unauthorized. */
export async function fetchPreviewTarget(
documentId: string,
chunkId?: string | null,
): Promise<PreviewTarget> {
const params = chunkId ? `?chunk_id=${encodeURIComponent(chunkId)}` : "";
const response = await authFetch(
`/api/rag/documents/${encodeURIComponent(documentId)}/preview-target${params}`,
);
return parseJsonOrThrow<PreviewTarget>(response);
}
/** Mint a short-lived URL that PDF.js can range-load without putting
* the user's bearer token in a query string. */
export async function fetchPreviewFileUrl(
documentId: string,
signal?: AbortSignal,
): Promise<PreviewFileUrl> {
const response = await authFetch(
`/api/rag/documents/${encodeURIComponent(documentId)}/file-url`,
signal ? { signal } : undefined,
);
const body = await parseJsonOrThrow<PreviewFileUrl>(response);
return {
...body,
url: apiUrl(body.url),
};
}
export async function backfillDocumentLocators(
documentId: string,
): Promise<LocatorBackfillResult> {
const response = await authFetch(
`/api/rag/documents/${encodeURIComponent(documentId)}/locators/backfill`,
{ method: "POST" },
);
return parseJsonOrThrow<LocatorBackfillResult>(response);
}
/** Download the original uploaded file as a Blob via `authFetch` (so
* the bearer token rides in the Authorization header never a query
* string). The caller (preview-store) creates and revokes the object
* URL so blob lifecycle stays in one place.
*
* `signal` lets the caller abort the fetch when the user switches
* documents mid-load. */
export async function fetchPreviewFileBlob(
documentId: string,
signal?: AbortSignal,
): Promise<Blob> {
const response = await authFetch(
`/api/rag/documents/${encodeURIComponent(documentId)}/file`,
signal ? { signal } : undefined,
);
if (!response.ok) {
const body = await response.json().catch(() => null);
throw new Error(parseErrorText(response.status, body));
}
return response.blob();
}

View file

@ -6,6 +6,7 @@ import { Button } from "@/components/ui/button";
import { cn } from "@/lib/utils";
import { Delete02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import type { KeyboardEvent, MouseEvent } from "react";
import type { RagDocument } from "../api/rag-api";
const STATUS_VARIANT: Record<
@ -27,20 +28,59 @@ function humanBytes(bytes: number): string {
export function DocumentRow({
doc,
onDelete,
onPreview,
rightSlot,
className,
}: {
doc: RagDocument;
onDelete?: () => void;
/** Fired when the row body (not the delete button) is clicked.
* Per decision Q9, callers should only pass this for completed
* documents. */
onPreview?: () => void;
rightSlot?: React.ReactNode;
className?: string;
}) {
const isPreviewable = !!onPreview;
const handleRowClick = (e: MouseEvent<HTMLDivElement>) => {
if (!onPreview) return;
// If a button/anchor/control was clicked (e.g. the delete icon),
// skip preview — let that handler win. Buttons inside this row
// additionally call stopPropagation, but this is defense in depth
// for any descendant Button that forgets to.
const target = e.target as HTMLElement | null;
const interactive = target?.closest("button, a, [role=button]");
if (interactive && interactive !== e.currentTarget) return;
onPreview();
};
const handleRowKey = (e: KeyboardEvent<HTMLDivElement>) => {
if (!onPreview) return;
if (e.target !== e.currentTarget) return; // ignore child key events
if (e.key === "Enter" || e.key === " ") {
e.preventDefault();
onPreview();
}
};
return (
<div
className={cn(
"flex items-center justify-between gap-3 rounded-md border border-border/60 px-3 py-2",
isPreviewable &&
"cursor-pointer outline-none focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50 hover:bg-muted/40",
className,
)}
{...(isPreviewable
? {
role: "button",
tabIndex: 0,
onClick: handleRowClick,
onKeyDown: handleRowKey,
"aria-label": `Open preview of ${doc.filename}`,
}
: {})}
>
<div className="flex min-w-0 flex-1 flex-col gap-1">
<div className="flex items-center gap-2">
@ -67,7 +107,12 @@ export function DocumentRow({
variant="ghost"
size="icon"
aria-label="Delete document"
onClick={onDelete}
onClick={(e) => {
// Stop propagation so the row's onClick (preview open)
// does not fire when the user is asking to delete.
e.stopPropagation();
onDelete();
}}
className="text-muted-foreground hover:text-destructive"
>
<HugeiconsIcon icon={Delete02Icon} size={16} />

View file

@ -7,6 +7,7 @@ import { Separator } from "@/components/ui/separator";
import { useState } from "react";
import type { KnowledgeBase } from "../api/rag-api";
import { useKBDocuments } from "../hooks/use-kb-documents";
import { usePreviewStore } from "../stores/preview-store";
import { DocumentRow } from "./document-row";
import { DocumentUploadDropzone } from "./document-upload-dropzone";
import { IngestionProgress } from "./ingestion-progress";
@ -14,10 +15,11 @@ import { KBReconfigureDialog } from "./kb-reconfigure-dialog";
export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) {
const { documents, loading, error, upload, remove } = useKBDocuments(kb.id);
const [activeJobsByDoc, setActiveJobsByDoc] = useState<Record<string, string>>(
{},
);
const [activeJobsByDoc, setActiveJobsByDoc] = useState<
Record<string, string>
>({});
const [reconfigureOpen, setReconfigureOpen] = useState(false);
const openPreview = usePreviewStore((s) => s.open);
const handleFiles = async (files: File[]) => {
for (const file of files) {
@ -70,9 +72,7 @@ export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) {
</span>
</div>
{error ? (
<div className="text-xs text-destructive">{error}</div>
) : null}
{error ? <div className="text-xs text-destructive">{error}</div> : null}
<ScrollArea className="flex-1">
<div className="flex flex-col gap-2 pr-2">
@ -92,6 +92,16 @@ export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) {
onDelete={() => {
void remove(doc.id);
}}
// Per decision Q9: only completed documents open a
// preview. Pending/running/failed rows degrade to
// non-interactive (no onPreview).
onPreview={
doc.status === "completed"
? () => {
void openPreview({ documentId: doc.id });
}
: undefined
}
rightSlot={
showProgress ? (
<IngestionProgress jobId={jobId} className="mt-1" />

View file

@ -0,0 +1,235 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import {
Sheet,
SheetContent,
SheetDescription,
SheetHeader,
SheetTitle,
} from "@/components/ui/sheet";
import { LoaderIcon, XIcon } from "lucide-react";
import {
type FC,
type ReactNode,
useEffect,
useSyncExternalStore,
} from "react";
import type { PreviewTarget } from "../api/rag-api";
import {
type PreviewLoadStatus,
isInlineBlobAllowed,
usePreviewStore,
} from "../stores/preview-store";
import { PreviewPdfView } from "./preview-pdf-view";
import { PreviewTextView } from "./preview-text-view";
import { PreviewUnavailable } from "./preview-unavailable";
interface PreviewPanelProps {
/** Whether the panel is currently being shown in its host slot.
* When the host hides the slot (e.g. the user closes both
* settings and preview from the chat header), we run the close
* side-effect so the blob URL is revoked. */
open: boolean;
disableDrawer?: boolean;
}
const LG_BREAKPOINT = 1024;
const MEDIA_QUERY = `(max-width: ${LG_BREAKPOINT - 1}px)`;
function getLgSnapshot(): boolean {
if (
typeof window === "undefined" ||
typeof window.matchMedia !== "function"
) {
return false;
}
return window.matchMedia(MEDIA_QUERY).matches;
}
function lgSubscribe(callback: () => void): () => void {
if (
typeof window === "undefined" ||
typeof window.matchMedia !== "function"
) {
return () => undefined;
}
const mql = window.matchMedia(MEDIA_QUERY);
mql.addEventListener("change", callback);
return () => mql.removeEventListener("change", callback);
}
function useIsViewportSqueezed(): boolean {
return useSyncExternalStore(lgSubscribe, getLgSnapshot, () => false);
}
interface PreviewBodyArgs {
error: string | null;
previewBlob: Blob | null;
previewFileUrl: string | null;
status: PreviewLoadStatus;
target: PreviewTarget | null;
}
function renderPreviewBody({
error,
previewBlob,
previewFileUrl,
status,
target,
}: PreviewBodyArgs): ReactNode {
if (status === "loading") {
return (
<div className="flex h-full items-center justify-center gap-2 text-xs text-muted-foreground">
<LoaderIcon className="size-3.5 animate-spin" />
Loading preview
</div>
);
}
if (status === "error") {
// Treat 404s as "document missing" so a stale citation reads
// like "no longer available" rather than a generic error.
const isMissing = (error ?? "").toLowerCase().includes("not found");
return (
<PreviewUnavailable
filename={target?.filename}
reason={error ?? "Preview unavailable."}
variant={isMissing ? "missing" : "error"}
/>
);
}
if (status === "ready" && target) {
const pdfFile = previewFileUrl ?? previewBlob;
if (
target.mediaKind === "pdf" &&
pdfFile &&
isInlineBlobAllowed(target.mediaKind)
) {
return <PreviewPdfView target={target} file={pdfFile} />;
}
// text / image / docx / html / unknown — all routed through
// text-view. text gets the snippet rendered inline; docx/html
// /unknown skip inline-render entirely (contracts §5.4 + Risk #3).
return <PreviewTextView target={target} />;
}
return null;
}
/** Body-only renderer. The host slot (desktop aside or mobile sheet)
* is owned by the host (chat-settings panel slot, kb-detail panel,
* etc.) this component is purely the content of the right slot. */
export const PreviewPanel: FC<PreviewPanelProps> = ({
open,
disableDrawer = false,
}) => {
const target = usePreviewStore((s) => s.target);
const previewBlob = usePreviewStore((s) => s.previewBlob);
const previewFileUrl = usePreviewStore((s) => s.previewFileUrl);
const status = usePreviewStore((s) => s.status);
const error = usePreviewStore((s) => s.error);
const close = usePreviewStore((s) => s.close);
const isSqueezed = useIsViewportSqueezed() && !disableDrawer;
// Unmount + visibility cleanup: when the panel is hidden or
// unmounted, revoke the live object URL (contracts §5.5).
useEffect(() => {
if (!open) {
close();
}
}, [open, close]);
useEffect(() => {
return () => {
// Component truly unmounting (e.g. navigation away). Cleanup
// anything still live.
usePreviewStore.getState().close();
};
}, []);
// Keyboard accessibility: ESC closes the preview.
useEffect(() => {
if (!open) {
return;
}
const onKey = (e: KeyboardEvent) => {
if (e.key === "Escape") {
e.preventDefault();
close();
}
};
document.addEventListener("keydown", onKey);
return () => document.removeEventListener("keydown", onKey);
}, [open, close]);
const body = renderPreviewBody({
error,
previewBlob,
previewFileUrl,
status,
target,
});
const renderedContent = (
<section
aria-label="Document preview"
className="flex h-full flex-col overflow-hidden bg-panel-surface/85 dark:bg-background/85 backdrop-blur-lg text-panel-surface-fg border border-border/40 shadow-lg menu-soft-surface"
>
<div className="flex items-center justify-between gap-2 border-b border-border/60 px-3 py-2 font-heading">
<div className="flex items-center gap-1.5">
<span
className="h-2 w-2 animate-pulse rounded-full bg-primary [--pulse-color:color-mix(in_oklab,var(--primary)_35%,transparent)]"
aria-hidden="true"
/>
<span className="text-sm font-semibold">Preview</span>
</div>
<Button
variant="ghost"
size="icon"
onClick={close}
aria-label="Close preview"
className="size-7"
>
<XIcon className="size-3.5" />
</Button>
</div>
<div className="min-h-0 flex-1 overflow-hidden">{body}</div>
</section>
);
if (isSqueezed) {
return (
<>
<div className="hidden" aria-hidden="true" />
<Sheet
open={open}
onOpenChange={(next) => {
if (!next) {
close();
}
}}
>
<SheetContent
side="right"
showCloseButton={false}
overlayClassName="bg-background/35 supports-backdrop-filter:backdrop-blur-[1px]"
className="preview-sheet-content p-0 font-heading data-[side=right]:w-full data-[side=right]:sm:max-w-md"
>
<SheetHeader className="sr-only">
<SheetTitle>Document preview</SheetTitle>
<SheetDescription>
Preview of the active document citation
</SheetDescription>
</SheetHeader>
<div className="flex h-full flex-col">{renderedContent}</div>
</SheetContent>
</Sheet>
</>
);
}
return renderedContent;
};

View file

@ -0,0 +1,579 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { cn } from "@/lib/utils";
import {
ChevronLeftIcon,
ChevronRightIcon,
CopyIcon,
LoaderIcon,
RotateCcwIcon,
SearchIcon,
ZoomInIcon,
ZoomOutIcon,
} from "lucide-react";
import {
type CSSProperties,
type FC,
useCallback,
useEffect,
useId,
useMemo,
useRef,
useState,
} from "react";
import { Document, Page, pdfjs } from "react-pdf";
import "react-pdf/dist/Page/AnnotationLayer.css";
import "react-pdf/dist/Page/TextLayer.css";
import type { PreviewPdfRegion, PreviewTarget } from "../api/rag-api";
import { PreviewUnavailable } from "./preview-unavailable";
// Configure pdfjs worker in the same module where react-pdf is used,
// per the react-pdf README. `import.meta.url` resolves to the JS bundle
// containing this module, and Vite (+ Tauri) rewrites the URL during
// build so the worker is co-located with the chunk.
pdfjs.GlobalWorkerOptions.workerSrc = new URL(
"pdfjs-dist/build/pdf.worker.min.mjs",
import.meta.url,
).toString();
type PreviewPdfFile = Blob | string;
interface PreviewPdfViewProps {
target: PreviewTarget;
file: PreviewPdfFile;
}
type LoadSuccess = { numPages: number };
type PdfLightThemeStyle = CSSProperties & Record<`--${string}`, string>;
const RESIZE_DEBOUNCE_MS = 100;
const MIN_PDF_WIDTH = 280;
// Body has p-2 (8px each side) + stable scrollbar gutter (~10px) + a tiny
// breathing margin so the page render doesn't kiss the scrollbar.
const PDF_BODY_GUTTER_PX = 28;
const PDF_THUMBNAIL_WIDTH = 64;
const PDF_LIGHT_THEME_STYLE: PdfLightThemeStyle = {
"--background": "oklch(1 0 0)",
"--foreground": "oklch(0.2686 0 0)",
"--card": "oklch(1 0 0)",
"--card-foreground": "oklch(0.1281 0.0179 169.2764)",
"--popover": "oklch(1 0 0)",
"--popover-foreground": "oklch(0.1281 0.0179 169.2764)",
"--primary": "#17b88b",
"--primary-foreground": "oklch(1 0 0)",
"--secondary": "oklch(0.9596 0.0275 167.8295)",
"--secondary-foreground": "oklch(0.2868 0.0649 159.9823)",
"--muted": "oklch(0.9702 0 0)",
"--muted-foreground": "oklch(0.5486 0 0)",
"--accent": "oklch(0.9596 0.0275 167.8295)",
"--accent-foreground": "oklch(0.2868 0.0649 159.9823)",
"--border": "oklch(0.9208 0.0101 164.8536)",
"--input": "oklch(0.9208 0.0101 164.8536)",
"--ring": "#17b88b",
colorScheme: "light",
};
function escapeHtml(value: string): string {
return value
.replace(/&/g, "&amp;")
.replace(/</g, "&lt;")
.replace(/>/g, "&gt;")
.replace(/"/g, "&quot;");
}
function markFirstMatch(text: string, needle: string): string | null {
const trimmed = needle.trim();
if (trimmed.length < 2) {
return null;
}
const lower = text.toLowerCase();
const start = lower.indexOf(trimmed.toLowerCase());
if (start < 0) {
return null;
}
const end = start + trimmed.length;
return `${escapeHtml(text.slice(0, start))}<mark>${escapeHtml(
text.slice(start, end),
)}</mark>${escapeHtml(text.slice(end))}`;
}
// Keep text-layer highlighting opt-in. Citation snippets render in the card
// below; using them here would mark common words across unrelated PDF text.
function highlightPdfText(text: string, searchTerm: string): string {
const trimmed = searchTerm.trim();
if (trimmed.length < 2) {
return escapeHtml(text);
}
const searchHit = markFirstMatch(text, trimmed);
if (searchHit) {
return searchHit;
}
return escapeHtml(text);
}
function regionIsOnPage(region: PreviewPdfRegion, pageNumber: number): boolean {
if (region.confidence !== "exact") {
return false;
}
if (region.pageNumber != null) {
return region.pageNumber === pageNumber;
}
return region.pageIndex === pageNumber - 1;
}
interface PdfThumbnailProps {
pageNumber: number;
active: boolean;
onSelect: (pageNumber: number) => void;
}
/** Lazy thumbnail rendered via IntersectionObserver only mounts the
* inner <Page> when scrolled into view (or close to it), so large PDFs
* stay responsive even when the rail caps at 80 buttons. */
const PdfThumbnail: FC<PdfThumbnailProps> = ({
pageNumber,
active,
onSelect,
}) => {
const buttonRef = useRef<HTMLButtonElement | null>(null);
const [shouldRender, setShouldRender] = useState(false);
useEffect(() => {
if (shouldRender) {
return;
}
const el = buttonRef.current;
if (!el || typeof IntersectionObserver === "undefined") {
// Fallback for jsdom / older browsers: render eagerly.
setShouldRender(true);
return;
}
const observer = new IntersectionObserver(
(entries) => {
for (const entry of entries) {
if (entry.isIntersecting) {
setShouldRender(true);
observer.disconnect();
return;
}
}
},
{ rootMargin: "320px" },
);
observer.observe(el);
return () => observer.disconnect();
}, [shouldRender]);
useEffect(() => {
const el = buttonRef.current;
if (!active || !el || typeof el.scrollIntoView !== "function") {
return;
}
el.scrollIntoView({ block: "nearest", behavior: "smooth" });
}, [active]);
return (
<button
ref={buttonRef}
type="button"
onClick={() => onSelect(pageNumber)}
aria-label={`Go to page ${pageNumber}`}
aria-current={active ? "page" : undefined}
className={cn(
"mb-1.5 flex w-full flex-col items-center gap-0.5 rounded-md p-1 outline-none transition-colors",
"focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-1",
active
? "bg-secondary/70 text-secondary-foreground"
: "hover:bg-muted/60",
)}
>
<div
className={cn(
"overflow-hidden rounded-sm border bg-white shadow-xs",
active
? "border-primary/70 ring-1 ring-primary/40"
: "border-border/60",
)}
style={{
width: PDF_THUMBNAIL_WIDTH,
minHeight: Math.round(PDF_THUMBNAIL_WIDTH * 1.3),
}}
>
{shouldRender ? (
<Page
pageNumber={pageNumber}
width={PDF_THUMBNAIL_WIDTH}
renderTextLayer={false}
renderAnnotationLayer={false}
loading={
<div
className="flex h-full w-full animate-pulse items-center justify-center bg-muted/40"
style={{
minHeight: Math.round(PDF_THUMBNAIL_WIDTH * 1.3),
}}
/>
}
error={
<div
className="flex h-full w-full items-center justify-center text-[8px] text-muted-foreground"
style={{
minHeight: Math.round(PDF_THUMBNAIL_WIDTH * 1.3),
}}
>
?
</div>
}
className="pointer-events-none [&_canvas]:!h-auto [&_canvas]:!w-full"
/>
) : null}
</div>
<span
className={cn(
"tabular-nums text-[10px]",
active ? "font-semibold" : "text-muted-foreground",
)}
>
{pageNumber}
</span>
</button>
);
};
export const PreviewPdfView: FC<PreviewPdfViewProps> = ({ target, file }) => {
const [numPages, setNumPages] = useState<number | null>(null);
const [pageNumber, setPageNumber] = useState<number>(target.targetPage ?? 1);
const [loadError, setLoadError] = useState<string | null>(null);
const [zoom, setZoom] = useState(1);
const [searchTerm, setSearchTerm] = useState("");
const [copied, setCopied] = useState(false);
const containerRef = useRef<HTMLDivElement | null>(null);
const lastMeasuredWidthRef = useRef<number | null>(null);
const lastResetKeyRef = useRef<string | null>(null);
const [width, setWidth] = useState<number | null>(null);
const searchInputId = useId();
const sourceKey =
typeof file === "string"
? file
: `${target.documentId}:${target.chunkId ?? ""}:${file.size}:${file.type}`;
const documentFile = useMemo(() => {
return typeof file === "string" ? { url: file } : file;
}, [file]);
const resetKey = `${sourceKey}:${target.targetPage ?? ""}`;
useEffect(() => {
if (lastResetKeyRef.current === resetKey) {
return;
}
lastResetKeyRef.current = resetKey;
setNumPages(null);
setLoadError(null);
setPageNumber(target.targetPage ?? 1);
setZoom(1);
setSearchTerm("");
setCopied(false);
}, [resetKey, target.targetPage]);
useEffect(() => {
const el = containerRef.current;
if (!el) {
return;
}
let timeoutId: number | null = null;
const updateWidth = () => {
const next = Math.max(MIN_PDF_WIDTH, el.clientWidth - PDF_BODY_GUTTER_PX);
if (lastMeasuredWidthRef.current === next) {
return;
}
lastMeasuredWidthRef.current = next;
setWidth(next);
};
updateWidth();
const observer = new ResizeObserver(() => {
if (timeoutId !== null) {
window.clearTimeout(timeoutId);
}
timeoutId = window.setTimeout(updateWidth, RESIZE_DEBOUNCE_MS);
});
observer.observe(el);
return () => {
observer.disconnect();
if (timeoutId !== null) {
window.clearTimeout(timeoutId);
}
};
}, []);
useEffect(() => {
if (!copied) {
return;
}
const id = window.setTimeout(() => setCopied(false), 1200);
return () => window.clearTimeout(id);
}, [copied]);
const handleLoadSuccess = useCallback(({ numPages }: LoadSuccess) => {
setNumPages(numPages);
setLoadError(null);
}, []);
const handleLoadError = useCallback((err: Error) => {
setLoadError(err.message || "Failed to load PDF");
}, []);
const goPrev = useCallback(() => {
setPageNumber((p) => Math.max(1, p - 1));
}, []);
const goNext = useCallback(() => {
setPageNumber((p) =>
numPages == null ? p + 1 : Math.min(numPages, p + 1),
);
}, [numPages]);
const textRenderer = useCallback(
({ str }: { str: string }) => highlightPdfText(str, searchTerm),
[searchTerm],
);
const currentRegions = useMemo(
() =>
(target.pdfRegions ?? []).filter((region) =>
regionIsOnPage(region, pageNumber),
),
[target.pdfRegions, pageNumber],
);
const visiblePageNumbers = useMemo(() => {
if (!numPages) {
return [];
}
const maxButtons = 80;
if (numPages <= maxButtons) {
return Array.from({ length: numPages }, (_, index) => index + 1);
}
const half = Math.floor(maxButtons / 2);
let start = Math.max(1, pageNumber - half);
const end = Math.min(numPages, start + maxButtons - 1);
start = Math.max(1, end - maxButtons + 1);
return Array.from({ length: end - start + 1 }, (_, index) => start + index);
}, [numPages, pageNumber]);
const pageWidth = width == null ? null : Math.round(width * zoom);
const excerptKey = `${sourceKey}:${target.chunkId ?? ""}:${
target.targetPage ?? ""
}:${pageNumber}`;
const copyExcerpt = useCallback(() => {
copyToClipboard(target.snippet ?? "").then(setCopied);
}, [target.snippet]);
if (loadError) {
return (
<PreviewUnavailable
filename={target.filename}
reason={loadError}
variant="error"
/>
);
}
return (
<div className="flex h-full min-h-0 flex-col">
<div className="flex flex-wrap items-center justify-between gap-2 border-b border-border/60 px-3 py-2 text-xs">
<span
className="min-w-0 flex-1 truncate font-semibold font-heading"
title={target.filename}
>
{target.filename}
</span>
<div className="flex shrink-0 items-center gap-2">
{/* Navigation Pill Group */}
<div className="flex items-center rounded-full border border-border/60 bg-muted/40 p-0.5 shadow-xs">
<Button
variant="ghost"
size="icon"
onClick={goPrev}
disabled={pageNumber <= 1}
aria-label="Previous page"
className="h-7 w-7 rounded-full hover:bg-background/80"
>
<ChevronLeftIcon className="size-3.5" />
</Button>
<span className="min-w-12 text-center tabular-nums text-[10px] font-medium text-muted-foreground">
{numPages == null
? `${pageNumber}/?`
: `${pageNumber}/${numPages}`}
</span>
<Button
variant="ghost"
size="icon"
onClick={goNext}
disabled={numPages != null && pageNumber >= numPages}
aria-label="Next page"
className="h-7 w-7 rounded-full hover:bg-background/80"
>
<ChevronRightIcon className="size-3.5" />
</Button>
</div>
{/* Zoom Pill Group */}
<div className="flex items-center rounded-full border border-border/60 bg-muted/40 p-0.5 shadow-xs">
<Button
variant="ghost"
size="icon"
onClick={() => setZoom((value) => Math.max(0.6, value - 0.1))}
aria-label="Zoom out"
className="h-7 w-7 rounded-full hover:bg-background/80"
>
<ZoomOutIcon className="size-3.5" />
</Button>
<Button
variant="ghost"
size="icon"
onClick={() => setZoom(1)}
aria-label="Reset zoom"
className="h-7 w-7 rounded-full hover:bg-background/80"
>
<RotateCcwIcon className="size-3.5" />
</Button>
<Button
variant="ghost"
size="icon"
onClick={() => setZoom((value) => Math.min(2.5, value + 0.1))}
aria-label="Zoom in"
className="h-7 w-7 rounded-full hover:bg-background/80"
>
<ZoomInIcon className="size-3.5" />
</Button>
</div>
{/* Copy Excerpt Pill Group */}
<div className="flex items-center rounded-full border border-border/60 bg-muted/40 p-0.5 shadow-xs">
<Button
variant="ghost"
size="icon"
onClick={copyExcerpt}
disabled={!target.snippet}
aria-label={
copied ? "Copied source excerpt" : "Copy source excerpt"
}
className="h-7 w-7 rounded-full hover:bg-background/80"
>
<CopyIcon className="size-3.5" />
</Button>
</div>
</div>
<label
htmlFor={searchInputId}
className="flex min-w-48 max-w-full flex-1 items-center gap-1 rounded-md border border-border/60 bg-background px-2"
>
<SearchIcon className="size-3.5 shrink-0 text-muted-foreground" />
<Input
id={searchInputId}
value={searchTerm}
onChange={(event) => setSearchTerm(event.target.value)}
placeholder="Search this PDF"
aria-label="Search this PDF"
className="h-7 border-0 bg-transparent px-0 text-xs shadow-none focus-visible:ring-0"
/>
</label>
</div>
{target.snippet ? (
<div
key={excerptKey}
className="m-2 rounded-lg border border-border/60 bg-muted/30 p-3 shadow-xs text-[11px] leading-relaxed text-foreground/80 transition-all duration-300 animate-in fade-in"
>
<p className="mb-1 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground/80">
Source Excerpt
{target.targetPage != null ? ` · Page ${target.targetPage}` : ""}
</p>
<p className="line-clamp-4 whitespace-pre-wrap font-sans text-muted-foreground">
{target.snippet}
</p>
</div>
) : null}
<Document
file={documentFile}
onLoadSuccess={handleLoadSuccess}
onLoadError={handleLoadError}
loading={
<div className="flex h-full items-center justify-center gap-2 text-xs text-muted-foreground">
<LoaderIcon className="size-3.5 animate-spin" />
Loading PDF...
</div>
}
error={
<PreviewUnavailable
filename={target.filename}
reason="The PDF could not be opened."
variant="error"
/>
}
className="flex min-h-0 flex-1"
>
<div className="preview-scrollbar w-[88px] shrink-0 overflow-y-auto border-r border-border/60 bg-muted/20 p-1.5">
{visiblePageNumbers.map((page) => (
<PdfThumbnail
key={page}
pageNumber={page}
active={page === pageNumber}
onSelect={setPageNumber}
/>
))}
</div>
<div
ref={containerRef}
className="preview-scrollbar flex-1 overflow-y-scroll overflow-x-auto bg-muted/20 p-2 [scrollbar-gutter:stable]"
>
<div
className="light [color-scheme:light] bg-white text-slate-900 rounded-md p-1 shadow-sm border border-border/30 [&_mark]:bg-primary/20 [&_mark]:text-slate-900 [&_mark]:ring-1 [&_mark]:ring-primary/60 [&_mark]:rounded-xs flex min-w-fit flex-col items-center"
style={PDF_LIGHT_THEME_STYLE}
>
{pageWidth != null ? (
<div
data-testid="pdf-main-page"
className="relative inline-block"
>
<Page
pageNumber={pageNumber}
width={pageWidth}
customTextRenderer={textRenderer}
renderTextLayer={true}
renderAnnotationLayer={false}
loading={
<div className="py-4 text-xs text-muted-foreground">
Rendering page...
</div>
}
className="shadow-sm"
/>
{currentRegions.map((region, index) => (
<div
key={`${region.pageIndex}-${region.x}-${region.y}-${index}`}
data-testid="pdf-region-highlight"
className="pointer-events-none absolute rounded-sm bg-primary/20 ring-1 ring-primary/60"
style={{
left: `${region.x * 100}%`,
top: `${region.y * 100}%`,
width: `${region.width * 100}%`,
height: `${region.height * 100}%`,
}}
/>
))}
</div>
) : null}
</div>
</div>
</Document>
</div>
);
};

View file

@ -0,0 +1,300 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import { authFetch } from "@/features/auth";
import { DownloadIcon, ExternalLinkIcon, FileTextIcon } from "lucide-react";
import { type FC, type ReactNode, useCallback } from "react";
import type { PreviewTarget } from "../api/rag-api";
import { isInlineBlobAllowed } from "../stores/preview-store";
interface PreviewTextViewProps {
target: PreviewTarget;
}
/** Fetch the original document bytes via authFetch so the bearer
* token rides in the Authorization header. `window.open(url)` and
* `<a download href=url>` cannot set custom headers, so handing
* them the raw `/file` URL gets a 401 (HTTPBearer-only backend
* see D1.3). */
async function fetchOriginalBlob(target: PreviewTarget): Promise<Blob> {
const response = await authFetch(
`/api/rag/documents/${encodeURIComponent(target.documentId)}/file`,
);
if (!response.ok) {
throw new Error(`Failed to fetch document (${response.status})`);
}
return response.blob();
}
function clickDownloadUrl(url: string, filename: string): void {
const anchor = document.createElement("a");
anchor.href = url;
anchor.download = filename;
anchor.style.display = "none";
document.body.appendChild(anchor);
anchor.click();
anchor.remove();
}
async function downloadOriginal(target: PreviewTarget): Promise<void> {
const blob = await fetchOriginalBlob(target);
const url = URL.createObjectURL(blob);
clickDownloadUrl(url, target.filename);
// Defer revocation so the browser's download pipeline gets the bytes
// before the URL goes away.
setTimeout(() => URL.revokeObjectURL(url), 0);
}
async function openOriginalInNewTab(target: PreviewTarget): Promise<void> {
// Defense in depth: refuse to create an inline blob URL for the
// unsafe types even if a future caller forgets the gate. The
// browser would render an html-blob as live HTML in the new tab,
// which is the Risk #3 / contracts §2.3 trip.
if (!isInlineBlobAllowed(target.mediaKind)) {
throw new Error(
`Inline open not allowed for mediaKind "${target.mediaKind}" — use Download instead.`,
);
}
const blob = await fetchOriginalBlob(target);
const url = URL.createObjectURL(blob);
window.open(url, "_blank", "noopener,noreferrer");
// Defer revocation so the new tab loads the bytes first.
setTimeout(() => URL.revokeObjectURL(url), 0);
}
/** Extracted-text / snippet preview used for:
* - `text` mediaKind (txt/md) the snippet is the only inline
* rendering we trust, and the original is one click away.
* - `docx`, `html`, `unknown` the original is NEVER rendered
* inline from an object URL (Risk #3); we show the cited chunk
* text plus a safe download/open action.
*
* When `chunk_id` was not supplied (document-row preview per
* contracts §1.3 + decision Q2), `snippet` is `null` and we show a
* metadata-only state instead of guessing a first chunk. */
interface MatchRange {
start: number;
end: number;
}
function findCharacterRange(
snippet: string,
target: PreviewTarget,
): MatchRange | null {
const { pageCharStart, pageCharEnd } = target;
if (pageCharStart !== null && pageCharEnd !== null) {
const start = Math.max(0, pageCharStart);
const end = Math.min(snippet.length, pageCharEnd);
if (start < end) {
return { start, end };
}
}
return null;
}
function findLineRange(
snippet: string,
target: PreviewTarget,
): MatchRange | null {
const { lineStart, lineEnd } = target;
if (lineStart !== null) {
const lines = snippet.split("\n");
const startLineIndex = Math.max(0, lineStart - 1);
const endLineIndex =
lineEnd !== null
? Math.min(lines.length - 1, lineEnd - 1)
: startLineIndex;
let charOffset = 0;
let startChar = -1;
let endChar = -1;
for (let i = 0; i < lines.length; i++) {
if (i === startLineIndex) {
startChar = charOffset;
}
charOffset += lines[i].length;
if (i === endLineIndex) {
endChar = charOffset;
break;
}
charOffset += 1; // for '\n'
}
if (startChar !== -1 && endChar !== -1 && startChar < endChar) {
return { start: startChar, end: endChar };
}
}
return null;
}
function findDensestLineRange(snippet: string): MatchRange | null {
const lines = snippet.split("\n");
let bestLineIndex = -1;
let maxAlphanumericCount = 0;
for (let i = 0; i < lines.length; i++) {
const alphanumericCount = lines[i].replace(/[^a-zA-Z0-9]/g, "").length;
if (alphanumericCount > maxAlphanumericCount) {
maxAlphanumericCount = alphanumericCount;
bestLineIndex = i;
}
}
if (bestLineIndex !== -1) {
let charOffset = 0;
for (let i = 0; i < bestLineIndex; i++) {
charOffset += lines[i].length + 1;
}
return { start: charOffset, end: charOffset + lines[bestLineIndex].length };
}
return null;
}
function findFuzzyMatch(
snippet: string,
target: PreviewTarget,
): MatchRange | null {
return (
findCharacterRange(snippet, target) ??
findLineRange(snippet, target) ??
findDensestLineRange(snippet)
);
}
const renderHighlightedSnippet = (
snippet: string,
target: PreviewTarget,
): ReactNode => {
const match = findFuzzyMatch(snippet, target);
if (!match) {
return snippet;
}
const before = snippet.slice(0, match.start);
const highlighted = snippet.slice(match.start, match.end);
const after = snippet.slice(match.end);
return (
<>
{before}
<mark className="rounded bg-primary/20 px-0.5 text-foreground ring-1 ring-primary/60">
{highlighted}
</mark>
{after}
</>
);
};
/** Extracted-text / snippet preview used for:
* - `text` mediaKind (txt/md) the snippet is the only inline
* rendering we trust, and the original is one click away.
* - `docx`, `html`, `unknown` the original is NEVER rendered
* inline from an object URL (Risk #3); we show the cited chunk
* text plus a safe download/open action.
*
* When `chunk_id` was not supplied (document-row preview per
* contracts §1.3 + decision Q2), `snippet` is `null` and we show a
* metadata-only state instead of guessing a first chunk. */
export const PreviewTextView: FC<PreviewTextViewProps> = ({ target }) => {
const snippet = target.snippet;
const hasSnippet = snippet !== null && snippet.trim().length > 0;
const hasLocator =
target.lineStart !== null ||
target.lineEnd !== null ||
target.pageCharStart !== null ||
target.pageCharEnd !== null;
// "Open original" creates a blob: URL of the original bytes and
// passes it to `window.open`. For `html` the new tab would render
// it as live HTML — exactly the Risk #3 / contracts §2.3 trip
// ("MUST refuse to create an inline object URL for mediaKind ==
// 'html' | 'docx' | 'unknown'"). For those types the only safe
// action is Download (backend already sets
// Content-Disposition: attachment for those Content-Types). The
// pdf/text/image allowlist is the same one the preview-store
// uses to decide whether to fetch the blob at all (§5.4). */
const canOpenInline = isInlineBlobAllowed(target.mediaKind);
const handleDownload = useCallback(() => {
downloadOriginal(target).catch(() => {
// best-effort; the user can retry the action.
});
}, [target]);
const handleOpenExternal = useCallback(() => {
// Re-fetch through authFetch and hand the new tab a blob URL.
// `window.open(rawApiUrl)` would send the request WITHOUT the
// Authorization header (window.open can't set custom headers)
// and the HTTPBearer-protected /file route would respond 401.
// See D1.3 finding.
openOriginalInNewTab(target).catch(() => {
// best-effort; the user can retry.
});
}, [target]);
return (
<div className="flex h-full flex-col gap-3 overflow-hidden p-4">
<div className="flex items-center gap-2 text-sm">
<FileTextIcon
className="size-4 shrink-0 text-muted-foreground"
aria-hidden={true}
/>
<span className="truncate font-medium" title={target.filename}>
{target.filename}
</span>
</div>
{target.targetPage != null ? (
<p className="text-xs text-muted-foreground">
Cited from page {target.targetPage}
</p>
) : null}
<div className="flex min-h-0 flex-1 flex-col gap-2 overflow-hidden rounded-md border border-border/60 bg-muted/30 p-3">
{hasSnippet ? (
<>
<p className="text-[10px] uppercase tracking-wider text-muted-foreground">
{hasLocator ? "Highlighted source excerpt" : "Source excerpt"}
</p>
{hasLocator ? (
<pre className="flex-1 overflow-auto whitespace-pre-wrap break-words text-xs leading-relaxed text-foreground/85">
{renderHighlightedSnippet(snippet, target)}
</pre>
) : (
<pre className="flex-1 overflow-auto whitespace-pre-wrap break-words text-xs leading-relaxed text-foreground/85">
{snippet}
</pre>
)}
</>
) : (
<p className="my-auto text-center text-xs text-muted-foreground">
{canOpenInline
? "No source excerpt — open the original to view this document."
: "No source excerpt — download the original to view this document."}
</p>
)}
</div>
<div className="flex gap-2">
{canOpenInline ? (
<Button
variant="outline"
size="sm"
onClick={handleOpenExternal}
className="flex-1"
>
<ExternalLinkIcon className="size-3.5" />
Open original
</Button>
) : null}
<Button
variant="outline"
size="sm"
onClick={handleDownload}
className="flex-1"
>
<DownloadIcon className="size-3.5" />
Download
</Button>
</div>
</div>
);
};

View file

@ -0,0 +1,46 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { AlertCircleIcon, FileXIcon } from "lucide-react";
import type { FC } from "react";
interface PreviewUnavailableProps {
/** Filename if known; "Document" otherwise. */
filename?: string;
/** One-line reason pulled from the backend's error body when
* available, otherwise a generic copy. */
reason: string;
/** "missing" deleted/404 case; "error" other failures. The icon
* + tone change so a stale citation reads as "no longer available"
* rather than a transient blip. */
variant?: "missing" | "error";
}
export const PreviewUnavailable: FC<PreviewUnavailableProps> = ({
filename,
reason,
variant = "error",
}) => {
const Icon = variant === "missing" ? FileXIcon : AlertCircleIcon;
const headline =
variant === "missing" ? "Document unavailable" : "Couldn't load preview";
return (
<output className="flex h-full flex-col items-center justify-center gap-3 px-6 text-center">
<Icon
className="size-10 text-muted-foreground"
strokeWidth={1.5}
aria-hidden={true}
/>
<div className="flex flex-col gap-1">
<p className="text-sm font-medium">{headline}</p>
{filename ? (
<p className="text-xs text-muted-foreground" title={filename}>
{filename}
</p>
) : null}
<p className="mt-1 max-w-xs text-xs text-muted-foreground">{reason}</p>
</div>
</output>
);
};

View file

@ -0,0 +1,158 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
type PointerEvent as ReactPointerEvent,
useCallback,
useEffect,
useRef,
useState,
} from "react";
interface UseResizableWidthOptions {
storageKey: string;
defaultWidth: number;
minWidth: number;
/** 0..1 fraction of viewport.innerWidth used as max width. Default 0.8. */
maxWidthFraction?: number;
/** Persist + listen for viewport-resize clamping only while true. */
enabled?: boolean;
}
interface UseResizableWidthResult {
width: number;
isResizing: boolean;
startResize: (event: ReactPointerEvent<HTMLElement>) => void;
adjustWidth: (delta: number) => void;
resetWidth: () => void;
}
function readStored(key: string, fallback: number): number {
if (typeof window === "undefined") {
return fallback;
}
try {
const raw = window.localStorage.getItem(key);
if (raw == null) {
return fallback;
}
const parsed = Number.parseInt(raw, 10);
return Number.isFinite(parsed) ? parsed : fallback;
} catch {
return fallback;
}
}
/** Drag-to-resize hook for a right-anchored panel. The handle sits on
* the panel's LEFT edge width grows as the pointer moves toward
* viewport x=0. Persists to localStorage and re-clamps on viewport
* changes so a wide panel cannot eclipse the host content. */
export function useResizablePanelWidth({
storageKey,
defaultWidth,
minWidth,
maxWidthFraction = 0.8,
enabled = true,
}: UseResizableWidthOptions): UseResizableWidthResult {
const [width, setWidth] = useState<number>(() =>
readStored(storageKey, defaultWidth),
);
const [isResizing, setIsResizing] = useState(false);
const rafRef = useRef<number | null>(null);
const clampWidth = useCallback(
(next: number): number => {
if (typeof window === "undefined") {
return Math.max(minWidth, next);
}
const max = Math.floor(window.innerWidth * maxWidthFraction);
return Math.max(minWidth, Math.min(max, next));
},
[minWidth, maxWidthFraction],
);
useEffect(() => {
if (!enabled || typeof window === "undefined") {
return;
}
try {
window.localStorage.setItem(storageKey, String(width));
} catch {
// localStorage may be unavailable (private mode, quota); persist
// is best-effort.
}
}, [width, storageKey, enabled]);
useEffect(() => {
if (typeof window === "undefined") {
return;
}
const onResize = () => {
setWidth((w) => clampWidth(w));
};
window.addEventListener("resize", onResize);
return () => window.removeEventListener("resize", onResize);
}, [clampWidth]);
const startResize = useCallback(
(event: ReactPointerEvent<HTMLElement>) => {
if (!enabled) {
return;
}
event.preventDefault();
const target = event.currentTarget;
const pointerId = event.pointerId;
try {
target.setPointerCapture(pointerId);
} catch {
// Pointer-capture isn't available everywhere (e.g. test envs).
}
setIsResizing(true);
document.body.style.cursor = "col-resize";
document.body.style.userSelect = "none";
const onMove = (e: PointerEvent) => {
if (rafRef.current !== null) {
cancelAnimationFrame(rafRef.current);
}
rafRef.current = requestAnimationFrame(() => {
setWidth(clampWidth(window.innerWidth - e.clientX));
});
};
const cleanup = (e: PointerEvent) => {
setIsResizing(false);
document.body.style.cursor = "";
document.body.style.userSelect = "";
try {
target.releasePointerCapture(e.pointerId);
} catch {
// Already released or not captured.
}
window.removeEventListener("pointermove", onMove);
window.removeEventListener("pointerup", cleanup);
window.removeEventListener("pointercancel", cleanup);
if (rafRef.current !== null) {
cancelAnimationFrame(rafRef.current);
rafRef.current = null;
}
};
window.addEventListener("pointermove", onMove);
window.addEventListener("pointerup", cleanup);
window.addEventListener("pointercancel", cleanup);
},
[enabled, clampWidth],
);
const adjustWidth = useCallback(
(delta: number) => {
setWidth((w) => clampWidth(w + delta));
},
[clampWidth],
);
const resetWidth = useCallback(() => {
setWidth(clampWidth(defaultWidth));
}, [clampWidth, defaultWidth]);
return { width, isResizing, startResize, adjustWidth, resetWidth };
}

View file

@ -0,0 +1,275 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { create } from "zustand";
import {
type PreviewMediaKind,
type PreviewTarget,
fetchPreviewFileBlob,
fetchPreviewFileUrl,
fetchPreviewTarget,
} from "../api/rag-api";
/** mediaKinds that may safely back an inline object URL (e.g. PDF.js
* worker, plain text, raster image). HTML, DOCX, and unknown are
* forced through the extracted-text fallback per contracts §5.4 +
* Risk #3 (unsafe HTML inline rendering). */
const INLINE_BLOB_ALLOWLIST: ReadonlySet<PreviewMediaKind> = new Set([
"pdf",
"text",
"image",
]);
export function isInlineBlobAllowed(mediaKind: PreviewMediaKind): boolean {
return INLINE_BLOB_ALLOWLIST.has(mediaKind);
}
/** What the panel should mount for the current target. Computed from
* `target.mediaKind` so the panel never has to re-derive it. */
export type PreviewLoadStatus = "idle" | "loading" | "ready" | "error";
export interface PreviewRequest {
/** Durable `rag_documents.id`. The only field required to open. */
documentId: string;
/** Durable `rag_chunks.id`. Optional absence means document-row
* preview (contracts §1.3 + decision Q2: snippet/targetPage stay
* null, no first-chunk fallback). */
backendChunkId?: string | null;
}
interface PreviewState {
/** Currently-open preview, or null when closed. */
target: PreviewTarget | null;
/** Object URL for the original file blob (PDF / text / image only).
* Null for docx / html / unknown (extracted-text fallback) and
* while the fetch is still in flight. */
previewBlobUrl: string | null;
/** Original fetched file blob for text/image fallback previews. PDFs
* prefer `previewFileUrl` so PDF.js can issue range requests. */
previewBlob: Blob | null;
/** Short-lived signed URL for PDF.js range requests. */
previewFileUrl: string | null;
previewFileUrlExpiresAt: number | null;
/** `target`-fetch + `blob`-fetch combined status. */
status: PreviewLoadStatus;
/** Last error message, if `status === "error"`. */
error: string | null;
/** Open key uniquely identifying the current request used by tests
* and by consumers that need to react to "the open call changed
* underneath me" (e.g. re-fetch after stale closure). */
openKey: number;
/** Open or replace the current preview. If a previous preview is
* open, its object URL is revoked and its in-flight fetch is
* aborted before the new request begins. */
open: (req: PreviewRequest) => Promise<void>;
/** Close the current preview. Revokes the object URL and aborts any
* in-flight fetch. Safe to call when nothing is open. */
close: () => void;
}
// State that can't live inside the zustand object without being
// part of the React render cycle. Kept module-scoped because the
// preview store is a singleton.
let activeAbortController: AbortController | null = null;
let activeBlobUrl: string | null = null;
let activeOpenKey = 0;
let restoreFocusElement: HTMLElement | null = null;
function revokeActiveBlobUrl(): void {
if (activeBlobUrl) {
URL.revokeObjectURL(activeBlobUrl);
activeBlobUrl = null;
}
}
function abortActive(): void {
if (activeAbortController) {
activeAbortController.abort();
activeAbortController = null;
}
}
export const usePreviewStore = create<PreviewState>((set) => ({
target: null,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "idle",
error: null,
openKey: 0,
async open(req) {
// Single-slot invariant (contracts §5.1): tear down whatever was
// there before assigning the new target. revoke → abort → reset.
revokeActiveBlobUrl();
abortActive();
activeOpenKey += 1;
const myKey = activeOpenKey;
const controller = new AbortController();
activeAbortController = controller;
const activeElement = document.activeElement;
restoreFocusElement =
activeElement instanceof HTMLElement ? activeElement : null;
set({
target: null,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "loading",
error: null,
openKey: myKey,
});
let target: PreviewTarget;
try {
target = await fetchPreviewTarget(
req.documentId,
req.backendChunkId ?? null,
);
} catch (err) {
if (myKey !== activeOpenKey) return; // superseded
if (activeAbortController === controller) activeAbortController = null;
set({
status: "error",
error: err instanceof Error ? err.message : String(err),
openKey: myKey,
});
return;
}
if (myKey !== activeOpenKey) {
// The user opened a different document while we were waiting.
return;
}
// For mediaKinds outside the allowlist (docx / html / unknown),
// skip the blob fetch entirely — the panel mounts the
// extracted-text fallback (contracts §5.4 + Risk #3).
if (!isInlineBlobAllowed(target.mediaKind)) {
if (activeAbortController === controller) activeAbortController = null;
set({
target,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "ready",
error: null,
openKey: myKey,
});
return;
}
if (target.mediaKind === "pdf") {
try {
const previewFile = await fetchPreviewFileUrl(
req.documentId,
controller.signal,
);
if (myKey !== activeOpenKey) return;
if (activeAbortController === controller) activeAbortController = null;
set({
target,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: previewFile.url,
previewFileUrlExpiresAt: previewFile.expiresAt,
status: "ready",
error: null,
openKey: myKey,
});
} catch (err) {
if (controller.signal.aborted || myKey !== activeOpenKey) return;
if (activeAbortController === controller) activeAbortController = null;
set({
target,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "error",
error: err instanceof Error ? err.message : String(err),
openKey: myKey,
});
}
return;
}
let blob: Blob;
try {
blob = await fetchPreviewFileBlob(req.documentId, controller.signal);
} catch (err) {
if (controller.signal.aborted || myKey !== activeOpenKey) return;
if (activeAbortController === controller) activeAbortController = null;
set({
target,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "error",
error: err instanceof Error ? err.message : String(err),
openKey: myKey,
});
return;
}
if (myKey !== activeOpenKey) {
// Superseded between target fetch and blob fetch — drop the bytes.
return;
}
const objectUrl = URL.createObjectURL(blob);
activeBlobUrl = objectUrl;
if (activeAbortController === controller) activeAbortController = null;
set({
target,
previewBlobUrl: objectUrl,
previewBlob: blob,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "ready",
error: null,
openKey: myKey,
});
},
close() {
revokeActiveBlobUrl();
abortActive();
activeOpenKey += 1; // poison any in-flight fetch that lands after this
const focusTarget = restoreFocusElement;
restoreFocusElement = null;
set({
target: null,
previewBlobUrl: null,
previewBlob: null,
previewFileUrl: null,
previewFileUrlExpiresAt: null,
status: "idle",
error: null,
openKey: activeOpenKey,
});
if (focusTarget?.isConnected) {
focusTarget.focus();
}
},
}));
/** Test-only inspector: returns whether the module-scoped blob URL is
* still live. Used by `preview-store.test.ts` to assert
* URL.revokeObjectURL was paired with URL.createObjectURL. */
export function __previewStoreInternals(): {
activeBlobUrl: string | null;
hasInflightController: boolean;
} {
return {
activeBlobUrl,
hasInflightController: activeAbortController !== null,
};
}

View file

@ -5,12 +5,36 @@ import { Separator } from "@/components/ui/separator";
import type { KnowledgeBase } from "@/features/rag/api/rag-api";
import { KBDetailPanel } from "@/features/rag/components/kb-detail-panel";
import { KBList } from "@/features/rag/components/kb-list";
import { PreviewPanel } from "@/features/rag/components/preview-panel";
import { RagDefaultsSection } from "@/features/rag/components/rag-defaults-section";
import { ThreadIndexList } from "@/features/rag/components/thread-index-list";
import { useState } from "react";
import { useResizablePanelWidth } from "@/features/rag/hooks/use-resizable-width";
import { usePreviewStore } from "@/features/rag/stores/preview-store";
import { cn } from "@/lib/utils";
import { type CSSProperties, useState } from "react";
export function KnowledgeBasesTab() {
const [selected, setSelected] = useState<KnowledgeBase | null>(null);
const previewTarget = usePreviewStore((s) => s.target);
const previewStatus = usePreviewStore((s) => s.status);
const previewActive =
previewTarget !== null ||
previewStatus === "loading" ||
previewStatus === "error";
const {
width: previewWidth,
isResizing: previewResizing,
startResize: startPreviewResize,
adjustWidth: adjustPreviewWidth,
resetWidth: resetPreviewWidth,
} = useResizablePanelWidth({
storageKey: "unsloth:kb-preview-panel-width",
defaultWidth: 560,
minWidth: 320,
maxWidthFraction: 0.7,
enabled: previewActive,
});
return (
<div className="flex h-full min-h-0 flex-col gap-4">
@ -36,6 +60,57 @@ export function KnowledgeBasesTab() {
</div>
)}
</div>
{previewActive ? (
<>
<button
type="button"
aria-label="Resize preview panel — drag, or use arrow keys"
aria-orientation="vertical"
onPointerDown={startPreviewResize}
onDoubleClick={resetPreviewWidth}
onKeyDown={(event) => {
if (event.key === "ArrowLeft") {
event.preventDefault();
adjustPreviewWidth(24);
} else if (event.key === "ArrowRight") {
event.preventDefault();
adjustPreviewWidth(-24);
} else if (event.key === "Home") {
event.preventDefault();
resetPreviewWidth();
}
}}
className={cn(
"group relative hidden w-1.5 cursor-col-resize touch-none select-none border-0 bg-transparent p-0 outline-none lg:block",
"before:absolute before:inset-y-0 before:left-1/2 before:-translate-x-1/2 before:w-px before:bg-border/70 before:transition-colors",
"hover:before:bg-primary/40 focus-visible:before:bg-primary/60",
previewResizing && "before:bg-primary/60",
)}
>
<span
aria-hidden="true"
className={cn(
"absolute left-1/2 top-1/2 h-10 w-[3px] -translate-x-1/2 -translate-y-1/2 rounded-full bg-border opacity-0 transition-opacity",
"group-hover:opacity-100 group-focus-visible:opacity-100",
previewResizing && "opacity-100",
)}
/>
</button>
<div
className={cn(
"w-0 shrink-0 overflow-hidden max-lg:hidden lg:w-[var(--preview-w)]",
!previewResizing && "transition-[width] duration-200 ease-out",
)}
style={
{
"--preview-w": `${previewWidth}px`,
} as CSSProperties
}
>
<PreviewPanel open={previewActive} />
</div>
</>
) : null}
</div>
<Separator />
<ThreadIndexList />

View file

@ -54,6 +54,7 @@
--duration-normal: 200ms;
/* Easing curves (Emil Kowalski) */
--ease-out-quad: cubic-bezier(0.25, 0.46, 0.45, 0.94);
--ease-out-quart: cubic-bezier(0.165, 0.84, 0.44, 1);
--ease-out-cubic: cubic-bezier(0.215, 0.61, 0.355, 1);
@ -1154,6 +1155,55 @@
background: oklch(0.72 0 0 / 0.25);
}
.preview-scrollbar {
scrollbar-width: thin;
scrollbar-color: oklch(0.5 0 0 / 0.25) transparent;
}
.preview-scrollbar::-webkit-scrollbar {
width: 10px;
height: 10px;
}
.preview-scrollbar::-webkit-scrollbar-track {
background: transparent;
}
.preview-scrollbar::-webkit-scrollbar-thumb {
background: oklch(0.5 0 0 / 0.22);
background-clip: padding-box;
border: 2px solid transparent;
border-radius: 9999px;
}
.preview-scrollbar:hover::-webkit-scrollbar-thumb {
background: oklch(0.5 0 0 / 0.38);
background-clip: padding-box;
border: 2px solid transparent;
}
.dark .preview-scrollbar {
scrollbar-color: oklch(0.72 0 0 / 0.25) transparent;
}
.dark .preview-scrollbar::-webkit-scrollbar-thumb {
background: oklch(0.72 0 0 / 0.25);
background-clip: padding-box;
border: 2px solid transparent;
}
.dark .preview-scrollbar:hover::-webkit-scrollbar-thumb {
background: oklch(0.78 0 0 / 0.42);
background-clip: padding-box;
border: 2px solid transparent;
}
.preview-sheet-content {
transition:
transform var(--duration-normal) var(--ease-out-quad),
opacity var(--duration-normal) var(--ease-out-quad);
}
/*---break---*/
@layer base {

View file

@ -0,0 +1 @@
import "@testing-library/jest-dom";

View file

@ -0,0 +1,19 @@
import { resolve } from "node:path";
import react from "@vitejs/plugin-react";
import { defineConfig } from "vitest/config";
export default defineConfig({
plugins: [react()],
test: {
environment: "jsdom",
globals: true,
setupFiles: ["./src/setupTests.ts"],
include: ["src/__tests__/**/*.{test,spec}.{ts,tsx}"],
exclude: ["node_modules", "dist"],
},
resolve: {
alias: {
"@": resolve(__dirname, "./src"),
},
},
});

View file

@ -0,0 +1,93 @@
"""Generate tests/fixtures/rag-preview/sample.pdf deterministically.
Run once: python tests/fixtures/rag-preview/make_fixture_pdf.py
Requires no third-party deps builds a minimal valid single-page PDF
using only stdlib so the fixture can be regenerated in any environment.
The output is committed alongside this script so tests load it directly.
"""
import os
import struct
import zlib
from pathlib import Path
OUTPUT = Path(__file__).parent / "sample.pdf"
def _compress(data: bytes) -> bytes:
return zlib.compress(data, level=9)
def _pdf() -> bytes:
# Minimal PDF 1.4 with one page, one text stream.
# Structure: header, catalog, pages, page, content stream, xref, trailer.
page_text = b"BT /F1 12 Tf 72 720 Td (RAG preview fixture - page 1) Tj ET"
compressed = _compress(page_text)
stream_len = len(compressed)
objects: list[bytes] = []
def obj(n: int, body: bytes) -> bytes:
return f"{n} 0 obj\n".encode() + body + b"\nendobj\n"
# 1: Catalog
objects.append(obj(1, b"<< /Type /Catalog /Pages 2 0 R >>"))
# 2: Pages
objects.append(obj(2, b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>"))
# 3: Page
objects.append(
obj(
3,
(
b"<< /Type /Page /Parent 2 0 R "
b"/MediaBox [0 0 612 792] "
b"/Contents 4 0 R "
b"/Resources << /Font << /F1 5 0 R >> >> >>"
),
)
)
# 4: Content stream
objects.append(
obj(
4,
(
f"<< /Length {stream_len} /Filter /FlateDecode >>".encode()
+ b"\nstream\n"
+ compressed
+ b"\nendstream"
),
)
)
# 5: Font
objects.append(
obj(
5,
b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
)
)
header = b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n"
body = b""
offsets: list[int] = []
for o in objects:
offsets.append(len(header) + len(body))
body += o
xref_offset = len(header) + len(body)
n = len(objects)
xref = f"xref\n0 {n + 1}\n".encode()
xref += b"0000000000 65535 f \n"
for off in offsets:
xref += f"{off:010d} 00000 n \n".encode()
trailer = (
f"trailer\n<< /Size {n + 1} /Root 1 0 R >>\n"
f"startxref\n{xref_offset}\n%%EOF\n"
).encode()
return header + body + xref + trailer
if __name__ == "__main__":
pdf_bytes = _pdf()
OUTPUT.write_bytes(pdf_bytes)
print(f"Written {len(pdf_bytes)} bytes to {OUTPUT}")

BIN
tests/fixtures/rag-preview/sample.pdf vendored Normal file

Binary file not shown.

8
tests/fixtures/rag-preview/sample.txt vendored Normal file
View file

@ -0,0 +1,8 @@
This is a test document for RAG preview fixtures.
Section 1: Introduction
The operating margin rose to 18.2% in Q3, driven by improved efficiency.
Section 2: Details
Additional supporting evidence and analysis is contained here.
Page 1 of 1.