diff --git a/app/api/kb/documents/[id]/route.ts b/app/api/kb/documents/[id]/route.ts index e68b0652..845c5dcf 100644 --- a/app/api/kb/documents/[id]/route.ts +++ b/app/api/kb/documents/[id]/route.ts @@ -1,12 +1,15 @@ import { NextResponse } from "next/server"; import { withAuth } from "@/lib/auth/with-auth"; +import { db } from "@/db/client"; import { deleteKbDocumentForUser, findKbChunksContentByDocumentId, findKbDocumentById, type KbChunkPreview, } from "@/lib/kb/queries"; +import { kbEntity, kbRelationship } from "@/lib/kb/schema"; +import { and, eq, sql } from "drizzle-orm"; // ponytail: Settings → KB → doc detail (right pane). Returns the // kb_document row + slim chunk preview (no 1536-dim embedding, no @@ -19,8 +22,26 @@ export const GET = withAuth<{ id: string }>(async (_req, { user, params }) => { return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); } - const chunks: KbChunkPreview[] = - doc.status === "success" ? await findKbChunksContentByDocumentId(user.id, doc.id) : []; + const [chunks, entityCount, relationshipCount] = await Promise.all([ + doc.status === "success" + ? findKbChunksContentByDocumentId(user.id, doc.id) + : Promise.resolve([] as KbChunkPreview[]), + // ponytail: parallel counts for the doc-detail badge so the + // "Indexed" check can require entity/relationship rows in addition + // to embeddings (hybrid-search third leg). Two separate queries + // (not a leftJoin) — joining two one-to-many tables would cartesian + // the rows and inflate both counts. + db + .select({ n: sql`count(*)::int` }) + .from(kbEntity) + .where(and(eq(kbEntity.documentId, doc.id), eq(kbEntity.userId, user.id))) + .then((rows) => rows[0]!.n), + db + .select({ n: sql`count(*)::int` }) + .from(kbRelationship) + .where(and(eq(kbRelationship.documentId, doc.id), eq(kbRelationship.userId, user.id))) + .then((rows) => rows[0]!.n), + ]); return NextResponse.json({ doc: { @@ -35,6 +56,8 @@ export const GET = withAuth<{ id: string }>(async (_req, { user, params }) => { pages: doc.pages, createdAt: doc.createdAt.toISOString(), updatedAt: doc.updatedAt.toISOString(), + entityCount, + relationshipCount, }, chunks, }); diff --git a/app/api/kb/documents/route.ts b/app/api/kb/documents/route.ts index 6c6d17a2..a6f85ec3 100644 --- a/app/api/kb/documents/route.ts +++ b/app/api/kb/documents/route.ts @@ -31,9 +31,12 @@ type GroupedDoc = { updatedAt: string; totalChunks?: number; successChunks?: number; + embeddingPendingChunks?: number; failedChunks?: number; pendingChunks?: number; parsingChunks?: number; + entityCount?: number; + relationshipCount?: number; totalPages?: number; failedPages?: number; pendingPages?: number; @@ -98,9 +101,12 @@ export const GET = withAuth(async (req: Request, { user }) => { updatedAt: d.updatedAt.toISOString(), totalChunks: d.totalChunks, successChunks: d.successChunks, + embeddingPendingChunks: d.embeddingPendingChunks, failedChunks: d.failedChunks, pendingChunks: d.pendingChunks, parsingChunks: d.parsingChunks, + entityCount: d.entityCount, + relationshipCount: d.relationshipCount, totalPages: d.totalPages, failedPages: d.failedPages, pendingPages: d.pendingPages, diff --git a/app/api/kb/folders/[id]/route.ts b/app/api/kb/folders/[id]/route.ts index 528f88ad..a429f4ff 100644 --- a/app/api/kb/folders/[id]/route.ts +++ b/app/api/kb/folders/[id]/route.ts @@ -84,16 +84,23 @@ export const PATCH = withAuth<{ id: string }>(async (req, { user, params }) => { // folder that still has docs. We surface that as 409 NON_EMPTY with // a doc count, so the UI can render "delete its 3 docs first" instead // of a generic FK-violation message. -export const DELETE = withAuth<{ id: string }>(async (_req, { user, params }) => { +export const DELETE = withAuth<{ id: string }>(async (req, { user, params }) => { const folder = await findKbFolderById(user.id, params.id); if (!folder) { return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); } - const docs = await listKbDocumentsByFolder(user.id, folder.id); - if (docs.length > 0) { - return NextResponse.json({ code: "NON_EMPTY", docCount: docs.length }, { status: 409 }); + + const url = new URL(req.url); + const deleteAll = url.searchParams.get("deleteAll") === "true"; + + if (!deleteAll) { + const docs = await listKbDocumentsByFolder(user.id, folder.id); + if (docs.length > 0) { + return NextResponse.json({ code: "NON_EMPTY", docCount: docs.length }, { status: 409 }); + } } - const deleted = await deleteKbFolderForUser(user.id, folder.id); + + const deleted = await deleteKbFolderForUser(user.id, folder.id, { deleteAllDocs: deleteAll }); if (!deleted) { return NextResponse.json({ code: "NOT_FOUND" }, { status: 404 }); } diff --git a/backend/node/kb/page-to-markdown-node.ts b/backend/node/kb/page-to-markdown-node.ts index 5c4c0440..c156d8cb 100644 --- a/backend/node/kb/page-to-markdown-node.ts +++ b/backend/node/kb/page-to-markdown-node.ts @@ -21,7 +21,18 @@ export const ocrPageSchema = z.object({ export async function pageToMarkdownNode( state: KbAgentStateShape, ): Promise> { - if (state.mode === "chunksOnly") { + if (state.mode === "chunksOnly" || state.mode === "retryFailedChunks") { + // ponytail: `chunksOnly` reuses the cached pages[].markdown verbatim. + // `retryFailedChunks` keeps doc.status='success' so all pages have + // markdown; only failed chunks (per-chunk UPDATE pre-applied by the + // route) need re-extraction. `retryFailed` is intentionally NOT in + // this list — its whole point is to re-OCR the failed pages, so this + // node has to run. The per-page OCR skip below (markdown.trim() > 0 + // AND no errorMessage) preserves successful pages while re-running + // OCR on any page with empty markdown OR a prior errorMessage, + // regardless of `status` — so `retryFailed`'s mixed-status + // pagesByDocId (some success, some failed) lands in the right lane + // per page. return { pagesByDocId: state.pagesByDocId, processedFiles: state.processedFiles, diff --git a/backend/node/kb/prepare-kb-data-node.ts b/backend/node/kb/prepare-kb-data-node.ts index c31ed18c..313cc706 100644 --- a/backend/node/kb/prepare-kb-data-node.ts +++ b/backend/node/kb/prepare-kb-data-node.ts @@ -44,7 +44,8 @@ export async function prepareKBDataNode( if (!userId) return makeError("user not provided"); const source = config?.configurable?.source ?? "chat"; - const parentMessageId = lastHumanMessageId(state.messages); + const parentMessageId = + config?.configurable?.parent_message_id ?? lastHumanMessageId(state.messages); if (!parentMessageId) { return makeError("kbAgent prepareKBDataNode: no HumanMessage to anchor parent_message_id on"); } diff --git a/backend/node/kb/split-file-page-node.ts b/backend/node/kb/split-file-page-node.ts index 13345c95..2af587d1 100644 --- a/backend/node/kb/split-file-page-node.ts +++ b/backend/node/kb/split-file-page-node.ts @@ -5,6 +5,17 @@ import { updateKbDocumentStatus } from "@/lib/kb/queries"; export async function splitFileToPageNode( state: KbAgentStateShape, ): Promise> { + if ( + state.mode === "chunksOnly" || + state.mode === "retryFailed" || + state.mode === "retryFailedChunks" + ) { + return { + pagesByDocId: state.pagesByDocId, + processedFiles: state.processedFiles, + }; + } + const newDocs = state.processedFiles.filter( (p) => p.pipelineStatus === "new" && p.docId !== null && p.r2Key !== null, ); diff --git a/backend/tool/kb/list-documents.ts b/backend/tool/kb/list-documents.ts index ec6aabf3..80ece68a 100644 --- a/backend/tool/kb/list-documents.ts +++ b/backend/tool/kb/list-documents.ts @@ -58,9 +58,12 @@ function toListDoc(row: { pendingPages?: number; totalChunks?: number; successChunks?: number; + embeddingPendingChunks?: number; failedChunks?: number; pendingChunks?: number; parsingChunks?: number; + entityCount?: number; + relationshipCount?: number; }): ListDocumentsDoc { const totalPages = row.totalPages ?? 0; const pending = row.pendingPages ?? 0; @@ -80,9 +83,12 @@ function toListDoc(row: { pendingPages: pending, totalChunks: row.totalChunks ?? 0, successChunks: row.successChunks ?? 0, + embeddingPendingChunks: row.embeddingPendingChunks ?? 0, failedChunks: row.failedChunks ?? 0, pendingChunks: row.pendingChunks ?? 0, parsingChunks: row.parsingChunks ?? 0, + entityCount: row.entityCount ?? 0, + relationshipCount: row.relationshipCount ?? 0, }; } diff --git a/backend/tool/kb/types.ts b/backend/tool/kb/types.ts index 4dfd8f9e..fff8a070 100644 --- a/backend/tool/kb/types.ts +++ b/backend/tool/kb/types.ts @@ -60,9 +60,12 @@ export type ListDocumentsDoc = { pendingPages: number; totalChunks: number; successChunks: number; + embeddingPendingChunks: number; failedChunks: number; pendingChunks: number; parsingChunks: number; + entityCount: number; + relationshipCount: number; }; export type ListDocumentsFolder = { diff --git a/components/assistant-ui/kb-mention.tsx b/components/assistant-ui/kb-mention.tsx index 4c96923d..7bb1b3b5 100644 --- a/components/assistant-ui/kb-mention.tsx +++ b/components/assistant-ui/kb-mention.tsx @@ -226,7 +226,7 @@ const KbMentionPopoverHeader = () => { if (scope.activeCategoryId) { return ( - + {"Back"} @@ -234,7 +234,7 @@ const KbMentionPopoverHeader = () => { } return ( -
+
Knowledge Base
); @@ -288,7 +288,7 @@ export const KbMentionPopover = memo(function KbMentionPopover() { {/* Directive must always be mounted so the popover can open when @ is typed */} diff --git a/components/settings/kb-view/dialogs.tsx b/components/settings/kb-view/dialogs.tsx index 5736f463..13e65a19 100644 --- a/components/settings/kb-view/dialogs.tsx +++ b/components/settings/kb-view/dialogs.tsx @@ -1,6 +1,6 @@ import { useCallback, useEffect, useState } from "react"; import { toast } from "sonner"; -import { Loader2 } from "lucide-react"; +import { Check, Copy, Loader2 } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Dialog, @@ -12,6 +12,7 @@ import { } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; +import { Switch } from "@/components/ui/switch"; import { cn } from "@/lib/utils"; import { TOAST_DESCRIPTION_CLASS, getModeInfo, type ReprocessMode } from "./helpers"; import { KbDocument, KbFolder } from "./types"; @@ -40,6 +41,7 @@ export function DocDeleteDialog({ descriptionClassName: TOAST_DESCRIPTION_CLASS, description: `「${doc.title}」was removed from this folder.`, }); + onOpenChange(false); onDeleted(); return; } @@ -53,7 +55,7 @@ export function DocDeleteDialog({ } finally { setSubmitting(false); } - }, [doc.id, onDeleted]); + }, [doc.id, doc.title, onDeleted, onOpenChange]); return ( -

- This permanently removes {doc.title} and: +

+ This permanently removes{" "} + {doc.title}{" "} + and:

  1. The document row + all parsed chunks (embeddings, BM25 index, entity graph)
    2. @@ -186,6 +190,7 @@ export function DocReprocessDialog({ descriptionClassName: TOAST_DESCRIPTION_CLASS, description: toastDesc, }); + onOpenChange(false); onReprocessed(); return; } @@ -216,7 +221,7 @@ export function DocReprocessDialog({ } finally { setSubmitting(false); } - }, [doc.id, doc.title, mode, onReprocessed]); + }, [doc.id, doc.title, mode, onOpenChange, onReprocessed]); return ( Reprocess document? - - {doc.title} - existing chunks are wiped before - re-running. Choose a reprocess mode. + + {doc.title} - + existing chunks are wiped before re-running. Choose a reprocess mode. @@ -339,66 +344,155 @@ export function FolderDeleteDialog({ }) { const [submitting, setSubmitting] = useState(false); const [error, setError] = useState(null); + const [deleteAll, setDeleteAll] = useState(false); + const [confirmText, setConfirmText] = useState(""); + const [copied, setCopied] = useState(false); useEffect(() => { setError(null); + setDeleteAll(false); + setConfirmText(""); + setCopied(false); }, [folder]); + const handleCopyClick = useCallback((e: React.MouseEvent) => { + e.preventDefault(); + e.stopPropagation(); + void navigator.clipboard.writeText("delete all").catch(() => {}); + toast.success("Copied 'delete all' to clipboard"); + setCopied(true); + setTimeout(() => setCopied(false), 1500); + }, []); + const submit = useCallback(async () => { if (!folder) return; setSubmitting(true); setError(null); try { - const res = await fetch(`/api/kb/folders/${folder.id}`, { method: "DELETE" }); + const url = `/api/kb/folders/${folder.id}${deleteAll ? "?deleteAll=true" : ""}`; + const res = await fetch(url, { method: "DELETE" }); if (res.status === 204) { toast.success("Folder deleted", { descriptionClassName: TOAST_DESCRIPTION_CLASS, - description: `「${folder.name}」was removed. Documents inside were kept.`, + description: deleteAll + ? `「${folder.name}」 and all its documents were permanently deleted.` + : `「${folder.name}」 was removed.`, }); + onOpenChange(false); onDeleted(); return; } if (res.status === 404) { toast.success("Folder deleted", { descriptionClassName: TOAST_DESCRIPTION_CLASS, - description: `「${folder.name}」was already removed.`, + description: `「${folder.name}」 was already removed.`, }); + onOpenChange(false); onDeleted(); return; } if (res.status === 409) { const body = (await res.json().catch(() => ({}))) as { docCount?: number }; setError( - `Folder still has ${body.docCount ?? "some"} document${body.docCount === 1 ? "" : "s"} — delete them first.`, + `Folder contains ${body.docCount ?? "some"} document${body.docCount === 1 ? "" : "s"}. Turn on the switch below to force clear and delete.`, ); return; } - setError(`Failed (${res.status})`); + setError(`Failed to delete folder (${res.status})`); } catch (e) { setError(e instanceof Error ? e.message : String(e)); } finally { setSubmitting(false); } - }, [folder, onDeleted]); + }, [folder, deleteAll, onDeleted, onOpenChange]); + + const isDeleteDisabled = + submitting || (deleteAll && confirmText.trim().toLowerCase() !== "delete all"); return ( { onOpenChange(o); - if (!o) setError(null); + if (!o) { + setError(null); + setDeleteAll(false); + setConfirmText(""); + } }} > Delete folder? - - This permanently removes the {folder?.name} folder. - Folders that still contain documents can't be deleted — empty the folder first. + + Are you sure you want to delete{" "} + {folder?.name}? - {error &&

    {error}

    } - + +
    + {error && ( +
    + {error} +
    + )} + +
    +
    + +

    + Also destroys all documents, chunks, and vector embeddings in this folder. Source + PDFs + rendered page images stay in R2 until the retention sweep. +

    +
    + { + setDeleteAll(c); + if (!c) setConfirmText(""); + }} + /> +
    + + {deleteAll && ( +
    + + setConfirmText(e.target.value)} + placeholder="delete all" + className="h-8 font-mono text-xs" + /> +
    + )} +
    + +
@@ -314,9 +317,12 @@ export function DocRow({ @@ -324,31 +330,23 @@ export function DocRow({
{actions}
- {previewOpen && ( - - )} - {deleteOpen && ( - { - setDeleteOpen(false); - void onRefresh(); - }} - /> - )} - {reprocessOpen && ( - { - setReprocessOpen(false); - void onRefresh(); - }} - /> - )} + + { + void onRefresh(); + }} + /> + { + void onRefresh(); + }} + /> ); } diff --git a/components/settings/kb-view/folder-sidebar.tsx b/components/settings/kb-view/folder-sidebar.tsx index 0fc44ff2..c10c00e1 100644 --- a/components/settings/kb-view/folder-sidebar.tsx +++ b/components/settings/kb-view/folder-sidebar.tsx @@ -169,9 +169,10 @@ export function FolderSidebar({ !o && setDeleteTarget(null)} + onOpenChange={(o) => { + if (!o) setDeleteTarget(null); + }} onDeleted={() => { - setDeleteTarget(null); void onRefresh(); }} /> @@ -179,9 +180,10 @@ export function FolderSidebar({ mode="edit" folder={editTarget} open={editTarget !== null} - onOpenChange={(o) => !o && setEditTarget(null)} + onOpenChange={(o) => { + if (!o) setEditTarget(null); + }} onSaved={() => { - setEditTarget(null); void onRefresh(); }} /> diff --git a/components/settings/kb-view/observability-popover.tsx b/components/settings/kb-view/observability-popover.tsx index 27076767..686eac3e 100644 --- a/components/settings/kb-view/observability-popover.tsx +++ b/components/settings/kb-view/observability-popover.tsx @@ -74,7 +74,6 @@ export function ObservabilityPopover({ docId }: { docId: string }) { }, [open, data, loading, docId]); const handleRunClick = (run: ObservabilityRun) => { - setOpen(false); // ponytail: each row carries its own threadId — chat uploads // open against the chat thread, standalone against the // docId-derived thread. parentMessageId is the synthetic diff --git a/components/settings/kb-view/types.ts b/components/settings/kb-view/types.ts index f427281e..1ccd4995 100644 --- a/components/settings/kb-view/types.ts +++ b/components/settings/kb-view/types.ts @@ -20,9 +20,12 @@ export type KbDocument = { updatedAt: string; totalChunks?: number; successChunks?: number; + embeddingPendingChunks?: number; failedChunks?: number; pendingChunks?: number; parsingChunks?: number; + entityCount?: number; + relationshipCount?: number; totalPages?: number; failedPages?: number; pendingPages?: number; @@ -43,6 +46,10 @@ export type KbChunkPreviewLocal = { themes: string[]; status: "pending" | "parsing" | "success" | "failed"; errorMessage: string | null; + // ponytail: mirrors lib/kb/queries.ts KbChunkPreview. Used by the + // doc-detail badge to detect status='success' chunks that haven't + // received a pgvector embedding yet ("Embedding" intermediate state). + hasEmbedding: boolean; }; export type KbDocDetail = { diff --git a/components/tool-ui/kb/list-documents-card.tsx b/components/tool-ui/kb/list-documents-card.tsx index 8089ebb0..18655592 100644 --- a/components/tool-ui/kb/list-documents-card.tsx +++ b/components/tool-ui/kb/list-documents-card.tsx @@ -127,9 +127,12 @@ function DocRow({ d }: { d: ListDocumentsFolder["documents"][number] }) { diff --git a/components/tool-ui/kb/status-badge.tsx b/components/tool-ui/kb/status-badge.tsx index e0e65cb1..ddbd1c5b 100644 --- a/components/tool-ui/kb/status-badge.tsx +++ b/components/tool-ui/kb/status-badge.tsx @@ -80,14 +80,14 @@ export function DocStatusBadge({ - {label} + {label} ); @@ -104,17 +104,32 @@ export function DocStatusBadge({ export function ChunksStatusBadge({ totalChunks, successChunks, + embeddingPendingChunks, failedChunks, pendingChunks, parsingChunks, + entityCount, + relationshipCount, docStatus, className, }: { totalChunks?: number; successChunks?: number; + // ponytail: chunks whose status='success' but embedding IS NULL — OCR + // is done but pgvector hasn't received the vector yet. Distinct from + // pending/parsing (chunks not yet OCR'd). Distinct from failed. + embeddingPendingChunks?: number; failedChunks?: number; pendingChunks?: number; parsingChunks?: number; + // ponytail: hybrid-search has three legs (BM25/tsv, pgvector/embedding, + // entity-tag overlap). "Indexed" requires at least one of entity or + // relationship rows to exist — otherwise the doc is queryable by vector + // but the entity-leg of hybrid search has nothing to score against. + // entity/relationship counts drive the third intermediate "Embedding" + // state (chunks all vectored but extraction still running). + entityCount?: number; + relationshipCount?: number; docStatus: KbStatus; className?: string; }) { @@ -127,12 +142,12 @@ export function ChunksStatusBadge({ - - Not Indexed + + Not Indexed ); } else if (docStatus === "pending" || docStatus === "parsing") { @@ -141,12 +156,12 @@ export function ChunksStatusBadge({ - - Pending + + Pending ); } else { @@ -162,19 +177,39 @@ export function ChunksStatusBadge({ - - No Chunks + + No Chunks ); } else { - const terminal = success + failed; + const embeddingPending = embeddingPendingChunks ?? 0; + const entities = entityCount ?? 0; + const relationships = relationshipCount ?? 0; + // ponytail: hybrid-search has three legs (BM25/tsv, pgvector, + // entity-tag overlap). "Indexed" is gated on vectors being + // complete (success === total) — the graph leg is supplementary + // and may legitimately produce zero rows for plain-text docs. + // Forcing `hasGraph` here would leave such docs stuck on + // "Embedding" forever with no terminal signal. The tooltip still + // surfaces the entity/relationship counts so the user sees the + // graph state explicitly. + const terminal = success + failed + embeddingPending; const isCompleted = success === total; - const isFailed = failed > 0 && terminal === total; - const isIndexing = !isCompleted && !isFailed; + // ponytail: "X failed" only fires when every non-failed chunk has + // actually finished OCR + embedding. If `embeddingPending > 0`, + // those chunks are still in flight (vector hasn't landed yet) — + // collapsing them into the "failed" bucket would render an + // "X failed" badge for a doc that's really "X failed, Y still + // embedding" and the user can't reconcile the count with the + // tooltip breakdown. + const isFailed = failed > 0 && embeddingPending === 0 && terminal === total; + const hasInflightChunks = pending > 0 || parsing > 0; + const hasInflightExtraction = entities === 0 && relationships === 0; + const isInProgress = !isCompleted && !isFailed; let variant: "success" | "destructive" | "muted" = "muted"; if (isFailed) { @@ -186,9 +221,31 @@ export function ChunksStatusBadge({ let label = "Indexed"; let iconElement = ; - if (isIndexing) { - label = `Indexing (${success}/${total})`; - tooltipText = `Indexing Status: ${success} chunks embedded, ${parsing} chunks in LLM extraction, ${pending} chunks queued, ${failed} chunks failed`; + if (isInProgress) { + // ponytail: three in-flight lanes collapse into two visible + // labels — "Extracting" (chunks still being created by OCR / + // chunk split) vs "Embedding" (chunks done but vector or + // graph-extract still running). The X/Y numerator is the count + // of chunks that have reached the *current lane's entry point*, + // not just the fully-embedded ones — otherwise the label/total + // don't sum to `total` and the user can't reconcile the badge + // with the tooltip breakdown. + if (hasInflightChunks) { + // OCR / chunk-split in flight. X = OCR'd chunks = success + + // embeddingPending. Y = total. Sums match the tooltip. + const ocrDone = success + embeddingPending; + label = `Extracting (${ocrDone}/${total})`; + tooltipText = `Indexing Status: ${success} chunks embedded, ${embeddingPending} chunks awaiting vector, ${parsing} chunks in extraction, ${pending} chunks queued, ${failed} chunks failed`; + } else { + // OCR done; vector and/or graph extraction in flight. X = + // embedded chunks. Y = total. + label = `Embedding (${success}/${total})`; + const vecBit = `${success} chunks embedded${embeddingPending > 0 ? `, ${embeddingPending} awaiting vector` : ""}`; + const extBit = hasInflightExtraction + ? `entity/relationship extraction in progress (${entities} entities, ${relationships} relationships so far)` + : ""; + tooltipText = `Indexing Status: ${vecBit}${extBit ? ", " + extBit : ""}`; + } const radius = 4.5; const circumference = 2 * Math.PI * radius; const pct = total > 0 ? success / total : 0; @@ -221,7 +278,13 @@ export function ChunksStatusBadge({ tooltipText = `Indexing Status: Vector database ingestion failed for ${failed} out of ${total} chunks`; iconElement = ; } else { - tooltipText = `Indexing Status: Successfully indexed all ${total} chunks to vector database`; + // ponytail: surface the graph counts in the tooltip even when + // the doc is "Indexed" — some docs legitimately produce zero + // entities/relationships (plain text, sparse pages) and the + // user should see that explicitly rather than wonder if the + // extraction is still running. + const graphBit = `${entities} entities, ${relationships} relationships`; + tooltipText = `Indexing Status: Successfully indexed all ${total} chunks to vector database (${graphBit})`; iconElement = ; } @@ -229,14 +292,14 @@ export function ChunksStatusBadge({ {iconElement} - {label} + {label} ); } diff --git a/components/tool-ui/kb/types.ts b/components/tool-ui/kb/types.ts index b14da6fd..804c4f4a 100644 --- a/components/tool-ui/kb/types.ts +++ b/components/tool-ui/kb/types.ts @@ -42,9 +42,12 @@ export type ListDocumentsDoc = { pendingPages: number; totalChunks: number; successChunks: number; + embeddingPendingChunks: number; failedChunks: number; pendingChunks: number; parsingChunks: number; + entityCount: number; + relationshipCount: number; }; export type ListDocumentsFolder = { diff --git a/docs/APIS.md b/docs/APIS.md index cb6c932c..bf8e0f2c 100644 --- a/docs/APIS.md +++ b/docs/APIS.md @@ -301,13 +301,16 @@ Rename a folder. Race-safe against concurrent renames — if two tabs race to th ### `DELETE /api/kb/folders/[id]` -Delete an empty folder. Postgres refuses via `ON DELETE RESTRICT` on `kb_document.folder_id` if any docs still belong to the folder; we surface that as 409 `NON_EMPTY` with a `docCount` so the UI can prompt "delete its N docs first". +Delete a folder. By default, Postgres refuses via `ON DELETE RESTRICT` on `kb_document.folder_id` if any docs still belong to the folder; we surface that as 409 `NON_EMPTY` with a `docCount` so the UI can prompt "delete its N docs first". -| | | -| ------------- | ---------------------------------------------------------------------- | -| Request body | (none) | -| 204 response | (empty body) | -| Failure codes | 401 `UNAUTHORIZED`. 404 `NOT_FOUND`. 409 `NON_EMPTY` (`{ docCount }`). | +With `?deleteAll=true`, the server cascades through every doc in the folder (each doc's `kb_chunk` rows, via `ON DELETE CASCADE`) before deleting the folder itself; the `attachment` rows / R2 objects backing those docs are NOT deleted (same retention as the single-doc `DELETE /api/kb/documents/[id]` semantics). The confirmation lives in the Settings → KB folder-delete dialog: a switch + a typed-confirmation of `delete all` is required before the request fires. + +| | | +| ------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Query params | `?deleteAll=true` → cascade-delete every doc in the folder (each doc's chunks, embeddings, attachments) before removing the folder. Omit for the empty-only default. | +| Request body | (none) | +| 204 response | (empty body) | +| Failure codes | 401 `UNAUTHORIZED`. 404 `NOT_FOUND`. 409 `NON_EMPTY` (`{ docCount }`) — only when `?deleteAll=true` is **not** set. | ### `GET /api/kb/documents` @@ -317,11 +320,13 @@ When `?mention=1` is set, returns a flat list of `status='success'` docs only When `?folderId=` is set, the sidebar still lists every folder the user owns, but only the targeted folder's `documents` array is populated. The other folders return `documents: []` so the Settings UI's "Live poll indicator" + scoped `anyInflight` only sees the folder the user is looking at. The `?folderId=` scope is ignored when `?mention=1` is set (the composer popover is always cross-folder). -| | | -| ------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Query params | `?mention=1` → flat `success`-only list for composer popover (omit for the default grouped response used by Settings). `?folderId=` → scope the doc payload to a single folder (Settings tab only). | -| 200 response | default: `{ groups: Array<{ folder: { id, name }, documents: Array<{ id, title, status, errorMessage, contentType, attachmentId, attachmentUrl, createdAt, updatedAt }> }> }`
`?mention=1`: `{ docs: Array<{ id, title, status }> }` | -| Failure codes | 401 `UNAUTHORIZED`. 500 `INTERNAL`. | +| | | +| ------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Query params | `?mention=1` → flat `success`-only list for composer popover (omit for the default grouped response used by Settings). `?folderId=` → scope the doc payload to a single folder (Settings tab only). | +| 200 response | default: `{ groups: Array<{ folder: { id, name, docCount }, documents: Array<{ id, title, status, errorMessage, contentType, attachmentId, attachmentUrl, createdAt, updatedAt, totalPages, failedPages, pendingPages, parsingPages, totalChunks, successChunks, embeddingPendingChunks, failedChunks, pendingChunks, parsingChunks, entityCount, relationshipCount }> }> }`
`?mention=1`: `{ folders: Array<{ id, name, docCount, docs: Array<{ id, title, status }> }> }` | +| Failure codes | 401 `UNAUTHORIZED`. 500 `INTERNAL`. | + +> **Indexing counts in the doc payload** — `successChunks` counts chunks with `status='success' AND embedding IS NOT NULL`. `embeddingPendingChunks` is chunks with `status='success' AND embedding IS NULL` (OCR done, vector pending). `entityCount` / `relationshipCount` are graph rows for the doc (hybrid-search third leg). The Settings → KB "Indexed" badge requires `successChunks === totalChunks AND (entityCount > 0 OR relationshipCount > 0)`; the chat-side `ChunksStatusBadge` enforces the same gate. ### `GET /api/kb/documents/[id]` @@ -329,10 +334,12 @@ Single-doc detail + slim chunk preview (text content only — no 1024-dim embedd > **Internal change, same contract.** Post issue #45 (audit §8 / migrations `0012_blue_steve_rogers.sql`, `0013_lowly_whirlwind.sql`, `0014_bright_norrin_radd.sql`), per-chunk `entities` / `themes` are no longer stored on `kb_chunk`. `entities` are JOINed server-side from `kb_entity` via `source_chunk_ids && ARRAY[chunkId]`; `themes` are JOINed from the flat `kb_theme` table on `chunk_id`; `relationships` from `kb_relationship`. Wire shape (`chunks[i].entities/relationships/themes`) is unchanged so the Settings → KB doc-detail UI reads them transparently. -| | | -| ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| 200 response | `{ doc: { id, title, status, errorMessage, contentType, attachmentId, folderId, contentHash, createdAt, updatedAt }, chunks: Array<{ ordinal: number, content: string, status, errorMessage, entities: Array<{name,type,description}>, relationships: Array<{source,target,relation,description}>, themes: string[] }> }` | -| Failure codes | 401 `UNAUTHORIZED`. 404 `NOT_FOUND`. | +| | | +| ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| 200 response | `{ doc: { id, title, status, errorMessage, contentType, attachmentId, folderId, contentHash, pages, createdAt, updatedAt, entityCount, relationshipCount }, chunks: Array<{ ordinal: number, content: string, status, errorMessage, hasEmbedding, entities: Array<{name,type,description}>, relationships: Array<{source,target,relation,description}>, themes: string[] }> }` | +| Failure codes | 401 `UNAUTHORIZED`. 404 `NOT_FOUND`. | + +> **`hasEmbedding`** — `true` iff `kb_chunk.embedding IS NOT NULL` (pgvector vector stored). A chunk with `status='success'` and `hasEmbedding=false` is in the "Embedding" intermediate state — OCR text was written but the vector hasn't been produced yet. The doc-level "Indexed" badge ignores chunks where `hasEmbedding=false` for its gate. ### `DELETE /api/kb/documents/[id]` diff --git a/docs/OBSERVABILITY.md b/docs/OBSERVABILITY.md index 4e03ad57..37b566f7 100644 --- a/docs/OBSERVABILITY.md +++ b/docs/OBSERVABILITY.md @@ -19,7 +19,7 @@ Each assistant message renders an icon-only `` inside its ` KB standalone ingestions (`fireIngestionRun` from `/settings/knowledge-base` upload + `/api/kb/documents/[id]/reprocess`) also surface in the same panel. The wiring path differs from the chat thread entry: - **`threadId` convention**: `docId.replace(/^d-/, "")` — strip the namespace prefix so the UUID portion satisfies LangGraph's `z.string().uuid()` validator and stays stable across reprocess. Every reprocess of the same doc lands on the same LangGraph thread; spans accumulate. -- **`parent_message_id` convention**: `fireIngestionRun` mints a fresh `messageId` per run and stamps it as `id` on the synthetic HumanMessage. `CapturingHandler.handleChainStart` reads it via `lastHumanMessageId` and tags every span in the run with `meta.parent_message_id = messageId`. The per-turn panel route (`/api/threads//observability/`) thus returns only that run's spans. +- **`parent_message_id` convention**: `fireIngestionRun` mints a fresh `messageId` per run and stamps it as `id` on the synthetic HumanMessage AND passes it to `langGraphClient.runs.create({ ... metadata: { parent_message_id: messageId } })` (top-level `metadata`, NOT `config.configurable` — see `lib/kb/ingest.ts` `fireIngestionRun`). LangGraph surfaces that on every LC callback's `metadata` for the run. `CapturingHandler.handleChainStart` reads `metadata.parent_message_id` first (per-run, immune to concurrent invokes) and falls back to `lastHumanMessageId(inputs.messages)` when the metadata key is absent (chat-path mainAgent invocations, which set pmid by other means). Every span in the run carries `meta.parent_message_id = messageId`, so the per-turn panel route (`/api/threads//observability/`) returns only that run's spans. - **Entry point on the UI**: each DocRow in `/settings/knowledge-base` has an `` icon (between RefreshCw and Search). Click opens a Popover ("Observability List") listing kbAgent re-runs for the doc — fetched from `GET /api/kb/documents/[id]/observability`, which reads the `kb_observability` table (no SDK call) and reshapes to `[{runId, threadId, parentMessageId, source, mode, createdAt}]` (newest first, limit 50). Initial `full`-mode uploads are NOT in the popover — the `kb_documents` row is the event for those; the popover shows chunksOnly / retryFailed / retryFailedChunks re-runs only. Each row is a button that calls `openSheet({ threadId: run.threadId, parentMessageId: run.parentMessageId })` against the same singleton `` used by chat messages. - **Settings page wiring**: `KbView` (`components/settings/kb-view/kb-view.tsx`) wraps its content in `` + ``, mirroring the chat thread tree. Without this provider, `useOpenObservabilitySheet()` throws "outside ObservabilitySheetProvider" on the KB page. - **Thread visibility in chat sidebar**: `threads.kind` (`'chat' | 'kb'`) is filtered at `lib/threads/queries.ts:listThreadsForUser` — `WHERE kind='chat'` keeps KB threads out of the user's chat sidebar. The settings page is the only UI entry into KB observability. @@ -159,6 +159,6 @@ curl -s -b cookies.txt -X DELETE \ - **`kind` includes `node` and `human`** — `node`: LangGraph's outer node-wrapper chain, distinct from LC `"chain"` wrappers. `human`: synthetic interrupt-wait gap; the transform layer maps `waiting` → `running` for the panel (which has no waiting state), while the DB keeps the precise value. - **`parent_span_id` is reconstructed from ns, not LC** — LC's `parent_run_id` reports root inside subgraphs, so the handler derives parents from `langgraph_checkpoint_ns` and rewrites `parent_span_id` before bulkInsert. - **Root chains are deduplicated by `meta.run_id`, not by `parent_span_id`** — under `streamSubgraphs: true` the LC inner `CompiledStateGraph` wrapper shares the outer wrapper's `meta.run_id` and `span_id === meta.run_id`, so `collectRootChains` keys by `run_id` to drop the duplicate. Two main invokes (regenerate + follow-up) emit two roots because their `run_id`s differ. -- **`parent_message_id` column backfill is in `bulkInsertSpans`** — interrupt-resume captures + cold-start threads (no `messages` in outermost inputs) leave `meta.parent_message_id` unset on the first span of the run. The next non-null pm id in the same thread backfills via the column. Pre-backfill or partially captured rows keep `parent_message_id IS NULL` and intentionally 404 on the per-turn detail endpoint (the wire's SpanData omits `parentMessageId` and the panel surfaces a "missing parent_message_id" error rather than sending a loose query). +- **`parent_message_id` resolution is per-call, not shared state** — concurrent invokes (e.g. two parallel kb uploads on different threads) previously clobbered each other via a process-wide `currentParentMessageId` instance field; the handler now resolves pmid inside `handleChainStart` from `metadata.parent_message_id` (LangGraph-configurable passthrough, set per-run via `runs.create`) with `lastHumanMessageId(inputs.messages)` as the fallback. `findLatestParentMessageId` (used by `bulkInsertSpans` backfill) calls `langGraphClient.threads.getState(threadId)` and parses `state.values.messages` via `lastHumanMessageId` — the column is **not** read for backfill. Pre-backfill or partially captured rows keep `parent_message_id IS NULL` and intentionally 404 on the per-turn detail endpoint (the wire's SpanData omits `parentMessageId` and the panel surfaces a "missing parent_message_id" error rather than sending a loose query). - **Redact instead of throw on forbidden fields** — avoids false-positive drops when user messages contain innocuous phrases like "api key". Auditable via `console.warn`. If zero-tolerance is required, change `redactForbidden` to throw. - **404 on cross-user thread access (not 401/403)** — deliberate, prevents enumeration. Implies ownership checks live at the route layer; the Sheet confirms by checking `useAuiState` threadId is real (not a `__LOCAL_*` placeholder) before mounting. diff --git a/docs/TOOLS.md b/docs/TOOLS.md index 4030a2a2..cd303575 100644 --- a/docs/TOOLS.md +++ b/docs/TOOLS.md @@ -107,7 +107,7 @@ Postgres `vector` extension; `list_documents` is unconditional. | Tool | Backend file | Frontend card | Notes | | ---------------- | ------------ | ----------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `search_KB` | `kb.ts` | `components/tool-ui/kb/search-kb-card.tsx` | Hybrid RRF (k=60) over three legs: BM25 (`tsv` GIN), pgvector cosine (`embedding` HNSW), and entity-tag overlap (`entities` GIN). Optional `folderId` / `documentId` filters add `WHERE` clauses inside the same SQL. Empty `rewriteQuery` falls back to ordinal-sorted chunks for the filtered scope (lets the LLM "summarize @doc.pdf" without a keyword); the dump cap is `KB_HYBRID_TOPK_DEFAULT` (8), not the previous 1000 — long-scope "summarize @doc" requests return at most 8 chunks, so doc-level summaries may be incomplete. When a Reranker model is registered, the candidate pool is widened to `max(50, topK * 5)` and the Reranker rescores; results below `KB_RERANK_MIN_SCORE` are filtered out before the final `topK` trim. **Partial Rerank caveat**: if the Reranker only returns scores for a subset of candidates, the unscored ones are discarded (no merge back with RRF scores), so a request for `topK` can return fewer than `topK`. Returns structured JSON: `{ content, documents[], empty }`. `content` embeds `[1] [2] [3]` markers; `documents[]` carries `chunkId`/`docId`/`rrfScore`/`legsHit` for UI. | -| `list_documents` | `kb.ts` | `components/tool-ui/kb/list-documents-card.tsx` | Paginated list of the user's KB docs. Filters: `folderId`, `status` (default `success`), `titleQuery` (ILIKE), `page` (default 1), `pageSize` (default 20, max 100). Strict filtering — no soft warnings. | +| `list_documents` | `kb.ts` | `components/tool-ui/kb/list-documents-card.tsx` | Paginated list of the user's KB docs. Filters: `folderId`, `status` (default `success`), `titleQuery` (ILIKE), `page` (default 1), `pageSize` (default 20, max 100). Strict filtering — no soft warnings. Each `documents[]` entry carries the same indexing counts the Settings → KB badges render (`totalPages` / `successPages` / `failedPages` / `parsingPages` / `pendingPages` + `totalChunks` / `successChunks` / `embeddingPendingChunks` / `failedChunks` / `pendingChunks` / `parsingChunks` + `entityCount` / `relationshipCount`) — single source of truth for the chat card and the Settings table. | ### ToolMessage shape @@ -120,6 +120,18 @@ Postgres `vector` extension; `list_documents` is unconditional. - `documents[]` — structured payload for the UI hover-cards: `{ chunkId, documentId, docTitle, pageNumbers, content, rrfScore, legsHit }`. +`list_documents` returns `{ content, folders, total, page, pageSize, empty }`: + +- `content` — terse LLM-facing summary listing folder names + doc titles + + statuses + error tails. No page/chunk counts (the model doesn't need + them to decide what to do next). +- `folders[]` — folder-grouped structured payload for the chat card. Each + `documents[]` entry mirrors the Settings → KB badge fields: + `{ id, title, status, errorMessage, createdAt, totalPages, successPages, +failedPages, parsingPages, pendingPages, totalChunks, successChunks, +embeddingPendingChunks, failedChunks, pendingChunks, parsingChunks, +entityCount, relationshipCount }`. + ### Lazy registration `search_KB` is always registered with the agent so the tool surface is diff --git a/lib/kb/queries.ts b/lib/kb/queries.ts index 3bd27960..3bac7a7b 100644 --- a/lib/kb/queries.ts +++ b/lib/kb/queries.ts @@ -293,10 +293,22 @@ export async function listKbDocumentsByFolder( export type KbDocumentWithAttachment = KbDocument & { attachmentUrl: string | null; totalChunks?: number; + // ponytail: successChunks now means "status='success' AND embedding IS NOT NULL" + // — the state that's actually queryable by hybrid search. Chunks that + // reached status='success' but are missing the embedding vector sit in + // embeddingPendingChunks and show as "Embedding" in the UI, not "Indexed". successChunks?: number; + embeddingPendingChunks?: number; failedChunks?: number; pendingChunks?: number; parsingChunks?: number; + // ponytail: kbEntity / kbRelationship row counts per doc. Hybrid search + // has three legs (BM25/tsv, pgvector/embedding, entity-tag overlap) — + // a doc is "Indexed" only when ALL THREE are populated. Vectors can + // arrive before entity/relationship extraction finishes; the UI shows + // "Embedding" until either count is > 0. + entityCount?: number; + relationshipCount?: number; totalPages?: number; failedPages?: number; pendingPages?: number; @@ -325,9 +337,19 @@ export async function listKbDocumentsByFolderWithAttachment( .select({ documentId: kbChunk.documentId, totalChunks: sql`count(${kbChunk.id})::int`.as("total_chunks"), - successChunks: sql`count(case when ${kbChunk.status} = 'success' then 1 end)::int`.as( - "success_chunks", - ), + // ponytail: successChunks = status='success' AND embedding IS NOT NULL. + // The status='success' flag alone doesn't mean the chunk is queryable; + // pgvector needs the vector. Chunks that finished the embed step but + // somehow ended up with a NULL vector (rare, but possible — embedder + // network blip, batch partial failure) land in embeddingPendingChunks. + successChunks: + sql`count(case when ${kbChunk.status} = 'success' and ${kbChunk.embedding} is not null then 1 end)::int`.as( + "success_chunks", + ), + embeddingPendingChunks: + sql`count(case when ${kbChunk.status} = 'success' and ${kbChunk.embedding} is null then 1 end)::int`.as( + "embedding_pending_chunks", + ), failedChunks: sql`count(case when ${kbChunk.status} = 'failed' then 1 end)::int`.as( "failed_chunks", ), @@ -342,15 +364,42 @@ export async function listKbDocumentsByFolderWithAttachment( .groupBy(kbChunk.documentId) .as("chunk_counts"); + // ponytail: separate subquery for the entity / relationship graph counts + // (kb_entity + kb_relationship rows grouped by document_id). Kept apart + // from chunksSubquery because they live in different tables — combining + // them in one GROUP BY would require a UNION, which loses the join + // semantics we want (each subquery's null side stays null, the outer + // coalesce handles docs that have zero rows in either table). + const entitySubquery = db + .select({ + documentId: kbEntity.documentId, + entityCount: sql`count(${kbEntity.id})::int`.as("entity_count"), + }) + .from(kbEntity) + .groupBy(kbEntity.documentId) + .as("entity_counts"); + + const relationshipSubquery = db + .select({ + documentId: kbRelationship.documentId, + relationshipCount: sql`count(${kbRelationship.id})::int`.as("relationship_count"), + }) + .from(kbRelationship) + .groupBy(kbRelationship.documentId) + .as("relationship_counts"); + const rows = await db .select({ doc: kbDocument, r2Key: attachments.r2Key, totalChunks: sql`coalesce(${chunksSubquery.totalChunks}, 0)::int`, successChunks: sql`coalesce(${chunksSubquery.successChunks}, 0)::int`, + embeddingPendingChunks: sql`coalesce(${chunksSubquery.embeddingPendingChunks}, 0)::int`, failedChunks: sql`coalesce(${chunksSubquery.failedChunks}, 0)::int`, pendingChunks: sql`coalesce(${chunksSubquery.pendingChunks}, 0)::int`, parsingChunks: sql`coalesce(${chunksSubquery.parsingChunks}, 0)::int`, + entityCount: sql`coalesce(${entitySubquery.entityCount}, 0)::int`, + relationshipCount: sql`coalesce(${relationshipSubquery.relationshipCount}, 0)::int`, totalPages: sql`coalesce(jsonb_array_length(${kbDocument.pages}), 0)::int`, // ponytail: each page carries an explicit `status` mirror of // kbChunkStatusEnum written by kb-agent's pageToMarkdownNode. @@ -375,6 +424,8 @@ export async function listKbDocumentsByFolderWithAttachment( .from(kbDocument) .leftJoin(attachments, eq(kbDocument.attachmentId, attachments.id)) .leftJoin(chunksSubquery, eq(kbDocument.id, chunksSubquery.documentId)) + .leftJoin(entitySubquery, eq(kbDocument.id, entitySubquery.documentId)) + .leftJoin(relationshipSubquery, eq(kbDocument.id, relationshipSubquery.documentId)) .where(and(eq(kbDocument.userId, userId), eq(kbDocument.folderId, folderId))) .orderBy(desc(kbDocument.createdAt)) .limit(limit); @@ -384,9 +435,12 @@ export async function listKbDocumentsByFolderWithAttachment( attachmentUrl: r.r2Key && base ? `${base}/${r.r2Key}` : null, totalChunks: r.totalChunks, successChunks: r.successChunks, + embeddingPendingChunks: r.embeddingPendingChunks, failedChunks: r.failedChunks, pendingChunks: r.pendingChunks, parsingChunks: r.parsingChunks, + entityCount: r.entityCount, + relationshipCount: r.relationshipCount, totalPages: r.totalPages, failedPages: r.failedPages, pendingPages: r.pendingPages, @@ -513,6 +567,12 @@ export type KbChunkPreview = { content: string; status: "pending" | "parsing" | "success" | "failed"; errorMessage: string | null; + // ponytail: whether pgvector has the vector for this chunk. status='success' + // doesn't imply this — embedder can drop a vector (network blip, batch + // partial failure). The doc-detail badge uses this to surface the + // "Embedding" intermediate state, same as the list endpoint's + // embeddingPendingChunks count. + hasEmbedding: boolean; entities: Array<{ name: string; type: string; description: string }>; relationships: Array<{ source: string; target: string; relation: string; description: string }>; themes: string[]; @@ -529,6 +589,7 @@ export async function findKbChunksContentByDocumentId( content: kbChunk.content, status: kbChunk.status, errorMessage: kbChunk.errorMessage, + hasEmbedding: sql`${kbChunk.embedding} is not null`.as("has_embedding"), }) .from(kbChunk) .innerJoin(kbDocument, eq(kbChunk.documentId, kbDocument.id)) @@ -628,6 +689,7 @@ export async function findKbChunksContentByDocumentId( content: c.content, status: c.status, errorMessage: c.errorMessage, + hasEmbedding: c.hasEmbedding, entities: entitiesByChunk.get(c.id) ?? [], relationships: relsByChunk.get(c.id) ?? [], themes: themesByChunk.get(c.id) ?? [], @@ -706,6 +768,7 @@ export async function findKbChunksByFolderId( content: kbChunk.content, status: kbChunk.status, errorMessage: kbChunk.errorMessage, + hasEmbedding: sql`${kbChunk.embedding} is not null`.as("has_embedding"), }) .from(kbChunk) .innerJoin(kbDocument, eq(kbChunk.documentId, kbDocument.id)) @@ -803,6 +866,7 @@ export async function findKbChunksByFolderId( content: c.content, status: c.status, errorMessage: c.errorMessage, + hasEmbedding: c.hasEmbedding, entities: entitiesByChunk.get(c.id) ?? [], relationships: relsByChunk.get(c.id) ?? [], themes: themesByChunk.get(c.id) ?? [], @@ -1432,6 +1496,16 @@ export async function deleteKbDocumentForUser( .where(and(eq(kbDocument.id, docId), eq(kbDocument.userId, userId))) .returning(); if (row) { + // ponytail: drop the attachment row too — same attachmentId can't be + // reused for another upload (sha-keyed dedup already prevents that), + // so leaving it behind would just be an orphan pointer to R2. R2 + // objects themselves stay (sha-keyed, reference-counted by future + // uploads with the same content — see CLAUDE.md / docs/ATTACHMENTS.md). + if (row.attachmentId) { + await tx + .delete(attachments) + .where(and(eq(attachments.id, row.attachmentId), eq(attachments.userId, userId))); + } const threadId = docId.replace(/^d-/, ""); await tx.delete(threads).where(and(eq(threads.id, threadId), eq(threads.kind, "kb"))); } @@ -1442,12 +1516,44 @@ export async function deleteKbDocumentForUser( export async function deleteKbFolderForUser( userId: string, folderId: string, + options?: { deleteAllDocs?: boolean }, ): Promise { - const [row] = await db - .delete(kbFolder) - .where(and(eq(kbFolder.id, folderId), eq(kbFolder.userId, userId))) - .returning(); - return row ?? null; + return db.transaction(async (tx) => { + if (options?.deleteAllDocs) { + const docs = await tx.query.kbDocument.findMany({ + where: and(eq(kbDocument.userId, userId), eq(kbDocument.folderId, folderId)), + columns: { id: true, attachmentId: true }, + }); + + for (const doc of docs) { + const threadId = doc.id.replace(/^d-/, ""); + // ponytail: thread cleanup precedes the kb_document delete on + // purpose. If a child cascade fails midway (an unenforced FK + // surfaces), the transaction rolls back — the threads stay + // attached to their docId strings and the next attempt re-runs + // safely. Reversing the order would orphan a thread pointing at + // a doc that's about to disappear. + await tx.delete(threads).where(and(eq(threads.id, threadId), eq(threads.kind, "kb"))); + // ponytail: drop the attachment row alongside the doc — see the + // note in deleteKbDocumentForUser. R2 objects stay. + if (doc.attachmentId) { + await tx + .delete(attachments) + .where(and(eq(attachments.id, doc.attachmentId), eq(attachments.userId, userId))); + } + } + + await tx + .delete(kbDocument) + .where(and(eq(kbDocument.userId, userId), eq(kbDocument.folderId, folderId))); + } + + const [row] = await tx + .delete(kbFolder) + .where(and(eq(kbFolder.id, folderId), eq(kbFolder.userId, userId))) + .returning(); + return row ?? null; + }); } export async function updateKbFolderNameForUser( diff --git a/lib/kb/source-kind.ts b/lib/kb/source-kind.ts index e94ae2fc..3c19c1cc 100644 --- a/lib/kb/source-kind.ts +++ b/lib/kb/source-kind.ts @@ -70,7 +70,10 @@ export function mimeShortLabel(mimeType: string): string { const office = OFFICE_LABEL[mt]; if (office) return office; const subtype = mt.split("/")[1] ?? ""; - return TYPE_LABEL[subtype] ?? (subtype.toUpperCase() || mimeType); + // ponytail: fall back to upper-cased subtype (e.g. "application/json" → "JSON") + // rather than the raw mime, so unmapped types still render readably. Empty + // subtype (malformed mime) falls through to the mime itself. + return TYPE_LABEL[subtype] ?? (subtype ? subtype.toUpperCase() : mimeType); } // ponytail: turn `R2_ALLOWED_CONTENT_TYPES` (comma-separated mimes) diff --git a/lib/observability/callback.ts b/lib/observability/callback.ts index 0cdb04ef..1f94d6de 100644 --- a/lib/observability/callback.ts +++ b/lib/observability/callback.ts @@ -250,14 +250,13 @@ export class CapturingHandler extends BaseCallbackHandler { // ticking until something else updates it. private actualParent = new Map(); - // ponytail: per-invoke last HumanMessage id. Set by the outermost - // handleChainStart from `inputs.messages`, inherited unchanged by - // nested spans in the same invoke. May be null when the run is a - // resume / regen / cold-start invoke with empty inputs — null is - // fine because bulkInsertSpans backfills from DB before INSERT, - // so spans in interrupted-or-recovered turns still tag with the - // thread's most recent non-null parent_message_id. - private currentParentMessageId: string | null = null; + // ponytail: parent_message_id is resolved per-call in start() from + // each span's own `partial.input.messages`. No instance-level cache, + // so concurrent invokes on the same handler (e.g. main chat + a + // background_agent dispatched mid-turn) can't clobber each other. + // Spans without messages in their input (inner LLM/tool nodes, or + // any background run that wasn't passed messages) get null here — + // bulkInsertSpans' state-fallback backfill fills them at INSERT time. constructor() { super(); @@ -268,25 +267,31 @@ export class CapturingHandler extends BaseCallbackHandler { chain: Serialized, inputs: Record, runId: string, - runType?: string, + _runType?: string, tags?: string[], metadata?: Record, runName?: string, parentRunId?: string, ) { - // ponytail: outermost call only — overwrite `currentParentMessageId` - // with the last HumanMessage id from `inputs.messages`. Null is OK - // for resume / regen / cold-start turns — bulkInsertSpans backfills - // from DB before INSERT. - if (!runType) { - this.currentParentMessageId = lastHumanMessageId((inputs as { messages?: unknown }).messages); - } + // ponytail: prefer metadata.parent_message_id (stamped by runs.create + // on kbAgent / background_agent dispatches) — it's per-run and won't + // be clobbered by concurrent invokes the way an instance field + // would. Fall back to parsing inputs.messages for the chat streaming + // path where metadata isn't set. + const metaPmid = metadata?.parent_message_id; + const fromMessages = lastHumanMessageId((inputs as { messages?: unknown }).messages); + const parentMessageId = + typeof metaPmid === "string" && metaPmid.length > 0 ? metaPmid : fromMessages; this.start(runId, parentRunId ?? null, { kind: "chain", name: runName ?? chain.id?.[chain.id.length - 1] ?? "chain", input: deepUnwrapLC(inputs), - meta: { ...metadata, ...(tags?.length ? { tags } : {}) }, + meta: { + ...metadata, + parent_message_id: parentMessageId, + ...(tags?.length ? { tags } : {}), + }, }); } @@ -528,18 +533,20 @@ export class CapturingHandler extends BaseCallbackHandler { // tags for nostream/etc., model name for header, and the TTFT. const meta = trimMeta(partial.meta); meta.time_to_first_token_ms ??= null; - meta.parent_message_id = this.currentParentMessageId; - - // ponytail: when `currentParentMessageId` is null at the outermost - // span (interrupt resume / regenerate / cold start — `inputs.messages` - // had no HumanMessage), the span gets stamped null here on purpose. - // Don't fire a DB lookup from this synchronous Start hook — the - // promise would resolve after handleChainEnd has already triggered - // bulkInsertSpans, and the `.then()` patch would land on a span - // that's been INSERTed (or about to be). The backfill in - // `bulkInsertSpans` reads the column at INSERT time and fills every - // null row in one pass — single source of truth, race-free by - // construction. + // ponytail: pmid fallback. handleChainStart pre-resolves from + // metadata.parent_message_id (per-run, set by runs.create) with a + // lastHumanMessageId(inputs.messages) fallback; the result is + // passed in via partial.meta.parent_message_id and we only run the + // message-array parse here when the caller didn't supply one. + // Inner LLM / tool / chain spans don't receive per-run metadata, so + // they end up with null — bulkInsertSpans' state-fallback backfill + // fills them at INSERT time. No instance-level cache: concurrent + // invokes can't clobber each other. + if (meta.parent_message_id == null) { + meta.parent_message_id = lastHumanMessageId( + (partial.input as { messages?: unknown } | null)?.messages, + ); + } this.spans.set(runId, { span_id: runId, diff --git a/lib/observability/queries.ts b/lib/observability/queries.ts index 0530af9a..7440a371 100644 --- a/lib/observability/queries.ts +++ b/lib/observability/queries.ts @@ -1,7 +1,9 @@ -import { and, desc, eq, isNotNull, lt, sql } from "drizzle-orm"; +import { and, eq, lt, sql } from "drizzle-orm"; import { db } from "@/db/client"; import { observabilitySpans, type NewObservabilitySpanRow } from "./schema"; import type { CapturedSpan } from "@/lib/observability/callback"; +import { langGraphClient } from "@/lib/langgraph/client"; +import { lastHumanMessageId } from "@/lib/langgraph/last-human-message-id"; // ponytail: FORBIDDEN regex flags fields that often hold secrets // (api_key, baseURL, organization, Bearer tokens). Spec FR-009 bans @@ -109,13 +111,16 @@ function toCapturedSpan(row: typeof observabilitySpans.$inferSelect): CapturedSp export async function bulkInsertSpans(spans: CapturedSpan[]): Promise { if (spans.length === 0) return 0; const rows = spans.map((s) => toRow(redactForbidden(s))); - // ponytail: parent_message_id backfill. Before INSERT, look up the - // column (not the meta key — the column is the canonical source) for - // any thread whose rows still have parent_message_id null. This is - // what the capture handler can't reach (race between outermost - // start()'s async DB lookup and bulkInsertSpans firing on End). - // Reads use the column so a row written by an earlier invoke but - // persisted without parent_message_id still drives the fill here. + // ponytail: parent_message_id backfill. The capture handler resolves + // pmid per-call from `metadata.parent_message_id` (LangGraph + // surfaces the `metadata` arg passed to `runs.create` on every LC + // callback) with `lastHumanMessageId(inputs.messages)` as the + // fallback — but `findLatestParentMessageId` is still useful for + // spans where both paths miss (cold-start threads, interrupt + // resumes whose capture preceded the metadata write). It calls + // `langGraphClient.threads.getState(threadId)` and parses + // `state.values.messages` via `lastHumanMessageId`; we never read + // the `observability_spans` column for the fill. await backfillParentMessageIds(rows); // ponytail: finalize any prior `status: "waiting"` interrupt spans on // the same thread before this batch's INSERTs. The handler used to @@ -124,6 +129,7 @@ export async function bulkInsertSpans(spans: CapturedSpan[]): Promise { // race-free and survives restart: any tool span arriving for the // thread implicitly closes the wait gap that opened on the previous // tool's interrupt. + await backfillWaitingInterruptSpans(rows); // ponytail: ON CONFLICT DO NOTHING makes the write idempotent — the // same runId can fire bulkInsertSpans twice (parent chain end + @@ -136,39 +142,69 @@ export async function bulkInsertSpans(spans: CapturedSpan[]): Promise { return inserted.length; } -// ponytail: latest non-null parent_message_id for the thread. Reads -// the column directly — meta->>'parent_message_id' jsonb path was -// dropped because the column is now the canonical store. Used both -// from CapturingHandler.start() (outermost-resume path) and from -// bulkInsertSpans' backfill (pre-INSERT gap-fill). +// ponytail: resolve the latest HumanMessage id from the canonical +// LangGraph thread state. State is the source of truth for "which +// message is current on this thread right now" — better than reading +// the spans column, where any prior turn could have written a row with +// a later started_at that pollutes earlier spans (e.g. kb-upload inner +// spans stealing the pmid of the next user message). Returns null on +// fetch error so bulkInsertSpans can proceed with null pmids rather +// than crash the run. export async function findLatestParentMessageId(threadId: string): Promise { - const [row] = await db - .select({ pmid: observabilitySpans.parentMessageId }) - .from(observabilitySpans) - .where( - and(eq(observabilitySpans.threadId, threadId), isNotNull(observabilitySpans.parentMessageId)), - ) - .orderBy(desc(observabilitySpans.startedAt)) - .limit(1); - const value = row?.pmid; - return typeof value === "string" && value.length > 0 ? value : null; + try { + const state = await langGraphClient.threads.getState(threadId); + const messages = (state as { values?: { messages?: unknown } } | null)?.values?.messages; + return lastHumanMessageId(messages); + } catch (err) { + console.warn( + `[observability] findLatestParentMessageId state fallback failed for thread ${threadId}:`, + err, + ); + return null; + } +} + +// ponytail: bound the state-fallback lookup so a stuck dev server +// (or a test env with no dev server) can't block the End hook past +// 500ms. The end-hook caller (handleChainEnd) awaits this whole +// INSERT before returning — a long hang here is directly visible as +// chat-turn latency. Returning null on timeout lets the row persist +// with `parent_message_id IS NULL`; the per-turn panel 404s for that +// span, which the design intentionally accepts (it surfaces "missing +// parent_message_id" instead of guessing). +function withTimeout(p: Promise, ms: number): Promise { + return Promise.race([p, new Promise((resolve) => setTimeout(() => resolve(null), ms))]); } -// ponytail: for each row whose parent_message_id column is null, look -// up the most recent non-null value on that thread and fill it in. -// One DB call per thread (deduplicated). Mutates the rows in place. +// ponytail: for each row whose parent_message_id column is null, ask +// findLatestParentMessageId (which now reads thread state via +// langGraphClient) for the latest HumanMessage id on that thread and +// stamp it in. One state call per thread (deduplicated). Mutates rows +// in place. async function backfillParentMessageIds(rows: NewObservabilitySpanRow[]): Promise { const missingThreads = new Set(); for (const r of rows) { - if (r.parentMessageId == null) missingThreads.add(r.threadId); + if (r.parentMessageId == null) { + missingThreads.add(r.threadId); + } } - for (const tid of missingThreads) { - const latest = await findLatestParentMessageId(tid); - if (!latest) continue; - for (const r of rows) { - if (r.threadId === tid && r.parentMessageId == null) { - r.parentMessageId = latest; - } + + // ponytail: parallelize the per-thread state lookups + 500ms cap each + // so a worst-case batch with N distinct missing threads adds at most + // ~500ms to the End hook (not N × RTT). Stale fetches resolve as null + // and the rows stay pmid-less; the per-turn panel 404s those rows + // by design. + const latestByThread = new Map(); + await Promise.all( + [...missingThreads].map(async (tid) => { + latestByThread.set(tid, await withTimeout(findLatestParentMessageId(tid), 500)); + }), + ); + + for (const r of rows) { + const latest = latestByThread.get(r.threadId); + if (r.parentMessageId == null && latest) { + r.parentMessageId = latest; } } } diff --git a/tests/api/kb/documents/[id].test.ts b/tests/api/kb/documents/[id].test.ts new file mode 100644 index 00000000..8a85c7fb --- /dev/null +++ b/tests/api/kb/documents/[id].test.ts @@ -0,0 +1,156 @@ +import "@/tests/helpers/session"; +import { afterAll, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import { eq } from "drizzle-orm"; + +import { db } from "@/db/client"; +import { attachments, kbDocument, kbEntity, kbFolder, kbRelationship, user } from "@/db/schema"; +import { setCurrentUser } from "@/tests/helpers/session"; + +import { GET } from "@/app/api/kb/documents/[id]/route"; + +// ponytail: GET /api/kb/documents/[id] — Settings → KB doc detail (right +// pane). The previous leftJoin(kbEntity, kbRelationship) cartesian- +// multiplied both counts whenever a doc had rows in both tables; the +// single-table subquery rewrite below fixes that. P1 from Greptile on PR +// #47. + +const USER_A = { id: "user-a", email: "a@x" }; +const USER_B = { id: "user-b", email: "b@x" }; + +beforeAll(() => { + process.env.LLM_KEY_ENCRYPTION_KEY ??= "a".repeat(64); +}); + +async function seedUser(u: { id: string; email: string }) { + await db + .insert(user) + .values({ id: u.id, email: u.email, name: "Test User" }) + .onConflictDoNothing(); +} + +async function seedFolder(userId: string, id: string, name: string) { + await db.insert(kbFolder).values({ id, userId, name }); +} + +async function seedDoc( + userId: string, + folderId: string, + id: string, + title: string, + status: "pending" | "parsing" | "success" | "failed" = "success", +) { + await db.insert(kbDocument).values({ + id, + userId, + folderId, + attachmentId: null, + title, + contentType: "application/pdf", + contentHash: `h-${id}`, + status, + }); +} + +async function seedEntity(userId: string, documentId: string, id: string, name: string) { + await db.insert(kbEntity).values({ + id, + userId, + documentId, + name, + type: "person", + description: "", + }); +} + +async function seedRelationship( + userId: string, + documentId: string, + id: string, + source: string, + target: string, +) { + await db.insert(kbRelationship).values({ + id, + userId, + documentId, + source, + target, + relation: "knows", + description: "", + }); +} + +beforeEach(async () => { + await db.delete(kbRelationship); + await db.delete(kbEntity); + await db.delete(kbDocument); + await db.delete(kbFolder); + await db.delete(attachments); + await db.delete(user).where(eq(user.id, USER_A.id)); + await db.delete(user).where(eq(user.id, USER_B.id)); + await seedUser(USER_A); + await seedUser(USER_B); + setCurrentUser(USER_A); +}); + +afterAll(() => { + setCurrentUser(null); +}); + +function newRequest(id: string): Request { + return new Request(`http://localhost/api/kb/documents/${id}`, { method: "GET" }); +} + +function ctxFor(id: string) { + return { params: Promise.resolve({ id }) }; +} + +describe("GET /api/kb/documents/[id] — graph counts", () => { + it("returns entityCount and relationshipCount = 0 for a doc with no graph rows", async () => { + setCurrentUser(USER_A); + await seedFolder(USER_A.id, "f-1", "Folder 1"); + await seedDoc(USER_A.id, "f-1", "d-1", "doc.pdf"); + + const res = await GET(newRequest("d-1"), ctxFor("d-1")); + expect(res.status).toBe(200); + const body = await res.json(); + expect(body.doc.entityCount).toBe(0); + expect(body.doc.relationshipCount).toBe(0); + }); + + it("returns correct counts when both entities and relationships exist (no cartesian inflation)", async () => { + setCurrentUser(USER_A); + await seedFolder(USER_A.id, "f-1", "Folder 1"); + await seedDoc(USER_A.id, "f-1", "d-1", "doc.pdf"); + + // 3 entities + 2 relationships — the old leftJoin would have + // cartesian'd into 6 rows and reported entityCount=6, relationshipCount=6. + await seedEntity(USER_A.id, "d-1", "e-1", "Alice"); + await seedEntity(USER_A.id, "d-1", "e-2", "Bob"); + await seedEntity(USER_A.id, "d-1", "e-3", "Carol"); + await seedRelationship(USER_A.id, "d-1", "r-1", "Alice", "Bob"); + await seedRelationship(USER_A.id, "d-1", "r-2", "Bob", "Carol"); + + const res = await GET(newRequest("d-1"), ctxFor("d-1")); + expect(res.status).toBe(200); + const body = await res.json(); + expect(body.doc.entityCount).toBe(3); + expect(body.doc.relationshipCount).toBe(2); + }); + + it("returns correct counts when only relationships exist (no entities)", async () => { + // Old leftJoin starting from kbEntity with no entity rows would + // return 0/0 (the from-table was empty). The single-table subquery + // rewrite counts each independently. + setCurrentUser(USER_A); + await seedFolder(USER_A.id, "f-1", "Folder 1"); + await seedDoc(USER_A.id, "f-1", "d-1", "doc.pdf"); + await seedRelationship(USER_A.id, "d-1", "r-1", "Alice", "Bob"); + + const res = await GET(newRequest("d-1"), ctxFor("d-1")); + expect(res.status).toBe(200); + const body = await res.json(); + expect(body.doc.entityCount).toBe(0); + expect(body.doc.relationshipCount).toBe(1); + }); +}); diff --git a/tests/api/kb/folder.test.ts b/tests/api/kb/folder.test.ts index 9cb664a5..a2e4391a 100644 --- a/tests/api/kb/folder.test.ts +++ b/tests/api/kb/folder.test.ts @@ -1,12 +1,21 @@ import "@/tests/helpers/session"; import { afterAll, beforeAll, beforeEach, describe, expect, it } from "vitest"; -import { eq } from "drizzle-orm"; +import { eq, sql } from "drizzle-orm"; import { db } from "@/db/client"; -import { attachments, kbChunk, kbDocument, kbFolder, user } from "@/db/schema"; +import { + attachments, + kbChunk, + kbDocument, + kbEntity, + kbFolder, + kbRelationship, + threads, + user, +} from "@/db/schema"; import { setCurrentUser } from "@/tests/helpers/session"; -import { GET } from "@/app/api/kb/folders/[id]/route"; +import { DELETE, GET } from "@/app/api/kb/folders/[id]/route"; // ponytail: GET /api/kb/folders/[id] returns the folder shape plus the // content + entities + status of every chunk across every doc in the @@ -26,12 +35,18 @@ async function seedFolderRow(userId: string, id: string, name = "Folder") { await db.insert(kbFolder).values({ id, userId, name }); } -async function seedDocRow(userId: string, folderId: string, id: string, title = "doc.pdf") { +async function seedDocRow( + userId: string, + folderId: string, + id: string, + title = "doc.pdf", + attachmentId: string | null = null, +) { await db.insert(kbDocument).values({ id, userId, folderId, - attachmentId: null, + attachmentId, title, contentType: "application/pdf", contentHash: `h-${id}`, @@ -39,14 +54,30 @@ async function seedDocRow(userId: string, folderId: string, id: string, title = }); } +async function seedAttachment(userId: string, id: string, r2Key: string) { + await db.insert(attachments).values({ + id, + userId, + r2Key, + name: `${id}.pdf`, + contentType: "application/pdf", + sizeBytes: 1024, + status: "uploaded", + sha256: `sha-${id}`, + }); +} + beforeAll(() => { process.env.LLM_KEY_ENCRYPTION_KEY ??= "a".repeat(64); }); beforeEach(async () => { + await db.delete(kbRelationship); + await db.delete(kbEntity); await db.delete(kbChunk); await db.delete(kbDocument); await db.delete(kbFolder); + await db.delete(threads); await db.delete(attachments); await db.delete(user).where(eq(user.id, USER_A.id)); await db.delete(user).where(eq(user.id, USER_B.id)); @@ -134,3 +165,129 @@ describe("GET /api/kb/folders/[id]", () => { expect(body.chunks).toHaveLength(2); }); }); + +describe("DELETE /api/kb/folders/[id]", () => { + it("returns 409 NON_EMPTY when folder has docs and deleteAll is false", async () => { + setCurrentUser(USER_A); + await seedFolderRow(USER_A.id, "f-1", "Folder 1"); + await seedDocRow(USER_A.id, "f-1", "d-1", "doc1.pdf"); + + const req = new Request("http://localhost/api/kb/folders/f-1", { method: "DELETE" }); + const res = await DELETE(req, ctxFor("f-1")); + expect(res.status).toBe(409); + const body = await res.json(); + expect(body.code).toBe("NON_EMPTY"); + expect(body.docCount).toBe(1); + }); + + it("deletes folder and all its documents when deleteAll=true", async () => { + setCurrentUser(USER_A); + await seedFolderRow(USER_A.id, "f-1", "Folder 1"); + await seedDocRow(USER_A.id, "f-1", "d-1", "doc1.pdf"); + await seedDocRow(USER_A.id, "f-1", "d-2", "doc2.pdf"); + + const req = new Request("http://localhost/api/kb/folders/f-1?deleteAll=true", { + method: "DELETE", + }); + const res = await DELETE(req, ctxFor("f-1")); + expect(res.status).toBe(204); + + const remainingDocs = await db.query.kbDocument.findMany({ + where: eq(kbDocument.folderId, "f-1"), + }); + expect(remainingDocs).toHaveLength(0); + + const folderRow = await db.query.kbFolder.findFirst({ + where: eq(kbFolder.id, "f-1"), + }); + expect(folderRow).toBeUndefined(); + }); + + it("deletes attachment rows for every doc when deleteAll=true (R2 objects stay)", async () => { + setCurrentUser(USER_A); + await seedAttachment(USER_A.id, "att-1", "u/user-a/upload/sha-att-1.pdf"); + await seedAttachment(USER_A.id, "att-2", "u/user-a/upload/sha-att-2.pdf"); + await seedFolderRow(USER_A.id, "f-1", "Folder 1"); + await seedDocRow(USER_A.id, "f-1", "d-1", "doc1.pdf", "att-1"); + await seedDocRow(USER_A.id, "f-1", "d-2", "doc2.pdf", "att-2"); + + const req = new Request("http://localhost/api/kb/folders/f-1?deleteAll=true", { + method: "DELETE", + }); + const res = await DELETE(req, ctxFor("f-1")); + expect(res.status).toBe(204); + + // Attachment rows for both docs are gone (no orphan pointers). + const remainingAttachments = await db.query.attachments.findMany({ + where: eq(attachments.userId, USER_A.id), + }); + expect(remainingAttachments).toHaveLength(0); + }); + + it("cascades: kb_chunk + kb_entity + kb_relationship rows gone, kb_thread rows deleted", async () => { + // ponytail: P1 follow-up from Greptile on PR #47. The earlier + // cascade test only verified the folder + kb_document rows went + // away; the turbo-button UX is one click away from "did the + // cascades actually run?" — add coverage for the per-doc kb_thread + // row + child rows on kb_chunk / kb_entity / kb_relationship. + setCurrentUser(USER_A); + await seedFolderRow(USER_A.id, "f-1", "Folder 1"); + await seedDocRow(USER_A.id, "f-1", "d-1", "doc1.pdf"); + + // Seed child rows directly (FK cascades do the work). + await db.insert(kbEntity).values({ + id: "e-1", + userId: USER_A.id, + documentId: "d-1", + name: "Alice", + type: "person", + description: "", + }); + await db.insert(kbRelationship).values({ + id: "r-1", + userId: USER_A.id, + documentId: "d-1", + source: "Alice", + target: "Bob", + relation: "knows", + description: "", + }); + // kb_chunk.tsv is a GENERATED ALWAYS column — raw INSERT skips the + // generated field that Drizzle insists on populating. + await db.execute( + sql`INSERT INTO kb_chunk (id, document_id, ordinal, content, status) VALUES ('c-1', 'd-1', 0, 'hello', 'success')`, + ); + // Per-doc kb thread (kind='kb') — deleteKbFolderForUser explicitly + // removes this in the cascade path; the single-doc DELETE does the + // same. Both are part of the FK CASCADE contract. + await db + .insert(threads) + .values({ id: "d-1".replace(/^d-/, ""), userId: USER_A.id, kind: "kb" }); + + const req = new Request("http://localhost/api/kb/folders/f-1?deleteAll=true", { + method: "DELETE", + }); + const res = await DELETE(req, ctxFor("f-1")); + expect(res.status).toBe(204); + + // FK cascades do the child-row work. + const remainingChunks = await db.query.kbChunk.findMany({ + where: eq(kbChunk.documentId, "d-1"), + }); + expect(remainingChunks).toHaveLength(0); + const remainingEntities = await db.query.kbEntity.findMany({ + where: eq(kbEntity.documentId, "d-1"), + }); + expect(remainingEntities).toHaveLength(0); + const remainingRelationships = await db.query.kbRelationship.findMany({ + where: eq(kbRelationship.documentId, "d-1"), + }); + expect(remainingRelationships).toHaveLength(0); + + // Explicit thread cleanup by deleteKbFolderForUser. + const remainingThreads = await db.query.threads.findMany({ + where: eq(threads.userId, USER_A.id), + }); + expect(remainingThreads).toHaveLength(0); + }); +}); diff --git a/tests/api/kb/reprocess.test.ts b/tests/api/kb/reprocess.test.ts index bc69085b..77f5291e 100644 --- a/tests/api/kb/reprocess.test.ts +++ b/tests/api/kb/reprocess.test.ts @@ -103,6 +103,7 @@ async function seedDoc(opts: { docId?: string; status?: "pending" | "parsing" | "success" | "failed"; errorMessage?: string | null; + pages?: unknown[]; }) { const docId = opts.docId ?? "d-1"; await db.insert(kbDocument).values({ @@ -115,6 +116,7 @@ async function seedDoc(opts: { contentHash: "h-1", status: opts.status ?? "success", errorMessage: opts.errorMessage ?? null, + pages: opts.pages ?? null, }); return docId; } @@ -407,4 +409,48 @@ describe("POST /api/kb/documents/[id]/reprocess", () => { // nothing and silently ignore. expect(mockRunsCreate).not.toHaveBeenCalled(); }); + + it("retryFailedChunks: prepareKBDataNode -> splitFileToPageNode -> pageToMarkdownNode pipeline skips OCR and preserves doc.status='success'", async () => { + const { prepareKBDataNode, splitFileToPageNode, pageToMarkdownNode } = + await import("@/backend/node/kb"); + await seedFolder(USER_A.id); + await seedAttachment(USER_A.id); + await seedDoc({ + userId: USER_A.id, + status: "success", + pages: [{ pageIndex: 0, imageUrl: "http://x", markdown: "hello" }], + }); + + const initialInput = { + messages: [{ id: "m-1", type: "human", content: "ingest" }], + }; + + const prepOut = await prepareKBDataNode(initialInput as never, { + configurable: { + userId: USER_A.id, + docId: "d-1", + mode: "retryFailedChunks", + source: "kb-reprocess", + parent_message_id: "m-1", + }, + }); + + expect(prepOut.status).toBe("success"); + expect(prepOut.pagesByDocId?.["d-1"]).toHaveLength(1); + + const mergedStateAfterPrep = { ...initialInput, ...prepOut }; + + const splitOut = await splitFileToPageNode(mergedStateAfterPrep as never); + expect(splitOut.pagesByDocId?.["d-1"]).toHaveLength(1); + + const mergedStateAfterSplit = { ...mergedStateAfterPrep, ...splitOut }; + + const ocrOut = await pageToMarkdownNode(mergedStateAfterSplit as never); + expect(ocrOut.pagesByDocId?.["d-1"]).toHaveLength(1); + + const docInDb = await db.query.kbDocument.findFirst({ + where: aAnd(eq(kbDocument.id, "d-1"), eq(kbDocument.userId, USER_A.id)), + }); + expect(docInDb?.status).toBe("success"); + }); }); diff --git a/tests/api/threads/observability.test.ts b/tests/api/threads/observability.test.ts index 77a2d976..dd29bbd7 100644 --- a/tests/api/threads/observability.test.ts +++ b/tests/api/threads/observability.test.ts @@ -1,7 +1,25 @@ import "@/tests/helpers/session"; -import { describe, it, expect, beforeAll, beforeEach, afterAll } from "vitest"; +import { describe, it, expect, beforeAll, beforeEach, afterAll, vi } from "vitest"; import { eq } from "drizzle-orm"; +// ponytail: spans with `langgraph_thread_id` but no `parent_message_id` +// trigger bulkInsertSpans → backfillParentMessageIds → +// findLatestParentMessageId → langGraphClient.threads.getState. In CI +// (no langgraph dev server) that call hangs. The bulkInsertSpans +// path has a 500ms timeout on the state fetch, but the test env wants +// a fast deterministic answer — stub the HTTP call so it returns null +// immediately and the row persists pmid-less (the panel 404s those by +// design). +vi.mock("@/lib/langgraph/client", () => ({ + langGraphClient: { + threads: { + getState: async () => ({ values: { messages: [] } }), + create: async () => ({}), + }, + runs: { create: async () => ({ run_id: "test" }) }, + }, +})); + import { db } from "@/db/client"; import { threads } from "@/lib/threads/schema"; import { observabilitySpans } from "@/lib/observability/schema"; diff --git a/tests/backend/node/kb/page-to-markdown-node.test.ts b/tests/backend/node/kb/page-to-markdown-node.test.ts new file mode 100644 index 00000000..a965783b --- /dev/null +++ b/tests/backend/node/kb/page-to-markdown-node.test.ts @@ -0,0 +1,139 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +const mockOcrInvoke = vi.fn(); +const mockUpdateKbDocumentStatus = vi.fn(); + +vi.mock("@/backend/model", () => ({ + getOcrModel: async () => ({ + withStructuredOutput: () => ({ + invoke: (...args: unknown[]) => mockOcrInvoke(...args), + }), + }), +})); + +vi.mock("@/lib/kb/queries", () => ({ + updateKbDocumentStatus: (...args: unknown[]) => mockUpdateKbDocumentStatus(...args), +})); + +import { pageToMarkdownNode } from "@/backend/node/kb/page-to-markdown-node"; + +// ponytail: P1 fix for Greptile on PR #47. The early-return guard +// incorrectly skipped OCR for retryFailed mode (only chunksOnly / +// retryFailedChunks should pass through). retryFailed re-OCRs the failed +// pages, so the OCR branch must run. + +beforeEach(() => { + mockOcrInvoke.mockReset(); + mockUpdateKbDocumentStatus.mockReset(); +}); + +describe("pageToMarkdownNode — retryFailed mode does NOT skip OCR", () => { + it("chunksOnly: returns cached pages by reference (OCR skipped)", async () => { + const pages = [{ pageIndex: 0, imageUrl: "http://x", markdown: "cached" }]; + const state = { + mode: "chunksOnly", + pagesByDocId: { "d-1": pages }, + processedFiles: [], + userId: "u-1", + }; + const out = await pageToMarkdownNode(state as never); + expect(out.pagesByDocId).toBe(state.pagesByDocId); + expect(mockOcrInvoke).not.toHaveBeenCalled(); + }); + + it("retryFailedChunks: returns cached pages by reference (OCR skipped)", async () => { + const pages = [{ pageIndex: 0, imageUrl: "http://x", markdown: "cached" }]; + const state = { + mode: "retryFailedChunks", + pagesByDocId: { "d-1": pages }, + processedFiles: [], + userId: "u-1", + }; + const out = await pageToMarkdownNode(state as never); + expect(out.pagesByDocId).toBe(state.pagesByDocId); + expect(mockOcrInvoke).not.toHaveBeenCalled(); + }); + + it("retryFailed: runs OCR — produces fresh pagesByDocId + invokes OCR model for failed pages", async () => { + mockOcrInvoke.mockResolvedValue({ markdown: "freshly OCR'd" }); + + const failedPage = { + pageIndex: 0, + imageUrl: "http://x", + markdown: "", + status: "failed", + errorMessage: "previous OCR timed out", + }; + const processedFile = { + docId: "d-1", + pipelineStatus: "new", + contentType: "application/pdf", + r2Key: "u/u-1/upload/sha.pdf", + title: "doc.pdf", + }; + const state = { + mode: "retryFailed", + pagesByDocId: { "d-1": [failedPage] }, + processedFiles: [processedFile], + userId: "u-1", + }; + + const out = await pageToMarkdownNode(state as never); + // Different reference → OCR branch ran (skip branch returns the + // exact same pagesByDocId reference). + expect(out.pagesByDocId).not.toBe(state.pagesByDocId); + // OCR was called for the failed page. + expect(mockOcrInvoke).toHaveBeenCalledTimes(1); + // The page markdown reflects the OCR result. + const outPages = out.pagesByDocId?.["d-1"] ?? []; + expect(outPages[0]?.markdown).toBe("freshly OCR'd"); + expect(outPages[0]?.status).toBe("success"); + }); + + it("retryFailed: per-page OCR skip preserves cached successful pages, re-runs only failed ones", async () => { + // ponytail: P2 follow-up from claude on PR #47 round 3. The per- + // page OCR skip preserves pages with non-empty markdown AND no + // errorMessage; a mixed pagesByDocId (some success, some failed) + // should land in the right lane per page. + mockOcrInvoke.mockResolvedValue({ markdown: "freshly OCR'd" }); + + const okPage = { + pageIndex: 0, + imageUrl: "http://x/0", + markdown: "page 0 already done", + status: "success", + }; + const failedPage = { + pageIndex: 1, + imageUrl: "http://x/1", + markdown: "", + status: "failed", + errorMessage: "previous OCR timed out", + }; + const processedFile = { + docId: "d-1", + pipelineStatus: "new", + contentType: "application/pdf", + r2Key: "u/u-1/upload/sha.pdf", + title: "doc.pdf", + }; + const state = { + mode: "retryFailed", + pagesByDocId: { "d-1": [okPage, failedPage] }, + processedFiles: [processedFile], + userId: "u-1", + }; + + const out = await pageToMarkdownNode(state as never); + const outPages = out.pagesByDocId?.["d-1"] ?? []; + + // OCR invoked only once — for the failed page. + expect(mockOcrInvoke).toHaveBeenCalledTimes(1); + // Cached page preserved verbatim. + expect(outPages[0]?.markdown).toBe("page 0 already done"); + expect(outPages[0]?.status).toBe("success"); + // Failed page re-OCRed. + expect(outPages[1]?.markdown).toBe("freshly OCR'd"); + expect(outPages[1]?.status).toBe("success"); + }); +}); diff --git a/tests/backend/observability/callback-collector.test.ts b/tests/backend/observability/callback-collector.test.ts index 5bfa1ce4..290ced07 100644 --- a/tests/backend/observability/callback-collector.test.ts +++ b/tests/backend/observability/callback-collector.test.ts @@ -128,7 +128,14 @@ describe("CapturingHandler — bulkInsert wiring", () => { }); describe("CapturingHandler — parent_message_id extraction", () => { - it("tags every span with the last HumanMessage id from the outermost chain's inputs", async () => { + it("outermost chain span resolves pmid from inputs.messages; inner LLM span stays null (backfill handles it)", async () => { + // ponytail: no shared instance state — the outermost chain span + // resolves pmid from its own inputs.messages via lastHumanMessageId. + // Inner LLM / tool spans don't receive the outermost messages in + // their input and don't get per-run metadata propagated by LangChain, + // so they get null at capture time. bulkInsertSpans' state-fallback + // backfill fills them at INSERT time from langGraphClient.threads. + // getState. const handler = makeHandler(); handler.handleChainStart( fakeSerialized("CompiledStateGraph"), @@ -155,14 +162,15 @@ describe("CapturingHandler — parent_message_id extraction", () => { handler.handleLLMEnd({ generations: [[]] } as never, "llm-1"); handler.handleChainEnd({ ok: true }, "outer"); - // ponytail: bulkInsert fires once per End hook. LLMEnd for llm-1 - // comes first → mock.calls[0] is [llm-1 span]. ChainEnd for outer - // comes second → mock.calls[1] is [outer span]. Both should carry - // the same parent_message_id since they're in the same invoke. + // bulkInsert fires once per End hook. LLMEnd for llm-1 first + // (mock.calls[0] = [llm-1 span]), then ChainEnd for outer + // (mock.calls[1] = [outer span]). const llmCall = (bulkInsertSpansMock.mock.calls[0] as unknown as [CapturedSpan[]])[0]; const chainCall = (bulkInsertSpansMock.mock.calls[1] as unknown as [CapturedSpan[]])[0]; expect(llmCall[0]?.span_id).toBe("llm-1"); - expect(llmCall[0]?.meta.parent_message_id).toBe("h-2"); + // inner span: no messages in its input, no per-run metadata → null + expect(llmCall[0]?.meta.parent_message_id).toBeNull(); + // outermost chain: parsed from inputs.messages expect(chainCall[0]?.span_id).toBe("outer"); expect(chainCall[0]?.meta.parent_message_id).toBe("h-2"); }); @@ -243,6 +251,66 @@ describe("CapturingHandler — parent_message_id extraction", () => { const callB = (bulkInsertSpansMock.mock.calls[1] as unknown as [CapturedSpan[]])[0]; expect(callB[0]?.meta.parent_message_id).toBe("h-B"); }); + + // ponytail: kb-upload / background_agent dispatches stamp + // metadata.parent_message_id via runs.create. The handler should + // honor that even if inputs.messages would resolve to a different + // value — metadata is the per-run trigger reference, the messages + // array may have been augmented by a later write. + it("metadata.parent_message_id wins over lastHumanMessageId(inputs.messages)", async () => { + const handler = makeHandler(); + handler.handleChainStart( + fakeSerialized("CompiledStateGraph"), + { + messages: [new HumanMessage({ content: "user-msg", id: "h-from-messages" })], + }, + "outer-bg", + undefined, + undefined, + { langgraph_thread_id: "t-bg", parent_message_id: "h-from-metadata" }, + ); + handler.handleChainEnd({ ok: true }, "outer-bg"); + const flushed = (bulkInsertSpansMock.mock.calls[0] as unknown as [CapturedSpan[]])[0]; + expect(flushed[0]?.meta.parent_message_id).toBe("h-from-metadata"); + }); + + // ponytail: the actual bug we're guarding against. Two concurrent + // invokes on the same handler instance — outer spans interleave + // before either ends. Pre-fix, the instance field + // `currentParentMessageId` would be overwritten by whichever + // handleChainStart ran last, polluting the other's spans. Post-fix + // there's no shared state; each handleChainStart resolves from its + // own inputs.messages. + it("concurrent invokes don't clobber each other's parent_message_id", async () => { + const handler = makeHandler(); + // invoke A: outermost h-A + handler.handleChainStart( + fakeSerialized("CompiledStateGraph"), + { messages: [new HumanMessage({ content: "a", id: "h-A" })] }, + "outer-A", + undefined, + undefined, + { langgraph_thread_id: "t-conc" }, + ); + // invoke B arrives before A ends: outermost h-B overwrites the + // instance field (pre-fix), but A's pmid is already stamped on its + // span via the per-call resolve. + handler.handleChainStart( + fakeSerialized("CompiledStateGraph"), + { messages: [new HumanMessage({ content: "b", id: "h-B" })] }, + "outer-B", + undefined, + undefined, + { langgraph_thread_id: "t-conc" }, + ); + handler.handleChainEnd({ ok: true }, "outer-A"); + handler.handleChainEnd({ ok: true }, "outer-B"); + + const callA = (bulkInsertSpansMock.mock.calls[0] as unknown as [CapturedSpan[]])[0]; + const callB = (bulkInsertSpansMock.mock.calls[1] as unknown as [CapturedSpan[]])[0]; + expect(callA[0]?.meta.parent_message_id).toBe("h-A"); // not h-B + expect(callB[0]?.meta.parent_message_id).toBe("h-B"); + }); }); describe("CapturingHandler — sync Start hooks stay sync", () => { diff --git a/tests/frontend/kb-list-documents-card.test.tsx b/tests/frontend/kb-list-documents-card.test.tsx index e46be27e..b38f7de9 100644 --- a/tests/frontend/kb-list-documents-card.test.tsx +++ b/tests/frontend/kb-list-documents-card.test.tsx @@ -40,9 +40,12 @@ function makeDoc(folderId: string, i: number): ListDocumentsDoc { pendingPages: 0, totalChunks: 24, successChunks: 24, + embeddingPendingChunks: 0, failedChunks: 0, pendingChunks: 0, parsingChunks: 0, + entityCount: 12, + relationshipCount: 5, }; } diff --git a/tests/integration/observability-wiring.test.ts b/tests/integration/observability-wiring.test.ts index c393f1ca..842a0c11 100644 --- a/tests/integration/observability-wiring.test.ts +++ b/tests/integration/observability-wiring.test.ts @@ -1,6 +1,20 @@ import "@/tests/helpers/session"; -import { describe, it, expect, beforeAll, afterAll } from "vitest"; +import { describe, it, expect, beforeAll, afterAll, vi } from "vitest"; import { eq } from "drizzle-orm"; + +// ponytail: see tests/api/threads/observability.test.ts for the why — +// bulkInsertSpans → backfillParentMessageIds hits +// langGraphClient.threads.getState, which hangs in CI without a dev +// server. Stub the HTTP call so the wiring test stays fast + deterministic. +vi.mock("@/lib/langgraph/client", () => ({ + langGraphClient: { + threads: { + getState: async () => ({ values: { messages: [] } }), + create: async () => ({}), + }, + runs: { create: async () => ({ run_id: "test" }) }, + }, +})); import { db } from "@/db/client"; import { threads } from "@/lib/threads/schema"; import { observabilitySpans } from "@/lib/observability/schema"; diff --git a/tests/lib/observability/queries.test.ts b/tests/lib/observability/queries.test.ts index 22c1f5db..bec89db6 100644 --- a/tests/lib/observability/queries.test.ts +++ b/tests/lib/observability/queries.test.ts @@ -1,5 +1,5 @@ import "@/tests/helpers/session"; -import { describe, it, expect, beforeEach } from "vitest"; +import { describe, it, expect, beforeEach, vi } from "vitest"; import { eq } from "drizzle-orm"; import { db } from "@/db/client"; @@ -13,6 +13,14 @@ import { deleteSpansOlderThan, findLatestParentMessageId, } from "@/lib/observability/queries"; + +const { mockGetState } = vi.hoisted(() => ({ + mockGetState: vi.fn<() => Promise<{ values?: { messages?: unknown } } | null>>(), +})); + +vi.mock("@/lib/langgraph/client", () => ({ + langGraphClient: { threads: { getState: mockGetState } }, +})); import { TEST_USER, ensureTestUser } from "@/tests/helpers/auth"; import type { CapturedSpan } from "@/lib/observability/callback"; @@ -208,6 +216,46 @@ describe("bulkInsertSpans", () => { const [fetched] = await getSpansByThreadId("t-roundtrip"); expect(fetched?.meta.parent_message_id).toBe("msg-from-column-only"); }); + + // ponytail: backfill asks langGraphClient for the latest HumanMessage + // id on the thread and stamps every null-pmid row in the batch with + // it. Rows that already carry a pmid are left alone. State is the + // canonical source — better than "latest in DB", which could be a + // later, unrelated turn. + it("backfill stamps every null-pmid row with the state's latest HumanMessage id", async () => { + await seedThread("t-x"); + mockGetState.mockResolvedValueOnce({ + values: { + messages: [ + { type: "human", id: "msg-1" }, + { type: "ai", id: "ai-1" }, + { type: "human", id: "msg-2" }, + ], + }, + }); + + await bulkInsertSpans([ + makeSpan({ + span_id: "wrapper", + meta: { langgraph_thread_id: "t-x", parent_message_id: "msg-1" }, // already set + }), + makeSpan({ + span_id: "inner-a", + meta: { langgraph_thread_id: "t-x" }, // null pmid — backfilled from state + }), + makeSpan({ + span_id: "inner-b", + meta: { langgraph_thread_id: "t-x" }, // null pmid — backfilled from state + }), + ]); + + const all = await getSpansByThreadId("t-x"); + // wrapper keeps its explicit pmid (not overwritten by state) + expect(all.find((s) => s.span_id === "wrapper")?.meta.parent_message_id).toBe("msg-1"); + // null-pmid rows get the state's latest HumanMessage id + expect(all.find((s) => s.span_id === "inner-a")?.meta.parent_message_id).toBe("msg-2"); + expect(all.find((s) => s.span_id === "inner-b")?.meta.parent_message_id).toBe("msg-2"); + }); }); describe("getSpansByThreadId", () => { @@ -340,48 +388,32 @@ describe("ON DELETE CASCADE — thread lifecycle", () => { }); }); -describe("findLatestParentMessageId (DB fallback)", () => { - // ponytail: this is what the collector's start() calls at the - // outermost span when currentParentMessageId is null. The most - // recent non-null parent_message_id on the thread is reused so a - // resumed / regen / cold-start invoke still tags with the right - // turn id. - it("returns the most recent parent_message_id for the thread", async () => { - await seedThread("t-fl"); - const base = Date.now(); - await bulkInsertSpans([ - makeSpan({ - span_id: "older", - started_at: base + 10, - meta: { langgraph_thread_id: "t-fl", parent_message_id: "msg-A" }, - }), - makeSpan({ - span_id: "newer", - started_at: base + 20, - meta: { langgraph_thread_id: "t-fl", parent_message_id: "msg-B" }, - }), - makeSpan({ - span_id: "no-pmid", - started_at: base + 30, - meta: { langgraph_thread_id: "t-fl" }, - }), - ]); - expect(await findLatestParentMessageId("t-fl")).toBe("msg-B"); +describe("findLatestParentMessageId (state fallback)", () => { + beforeEach(() => { + mockGetState.mockReset(); }); - it("returns null when the thread has no prior parent_message_id", async () => { - await seedThread("t-empty"); - await bulkInsertSpans([ - makeSpan({ - span_id: "no-pmid-1", - meta: { langgraph_thread_id: "t-empty" }, - }), - ]); - expect(await findLatestParentMessageId("t-empty")).toBeNull(); + it("returns the last HumanMessage id from langGraphClient thread state", async () => { + mockGetState.mockResolvedValueOnce({ + values: { + messages: [ + { type: "human", id: "msg-X" }, + { type: "ai", id: "ai-1" }, + { type: "human", id: "msg-Y" }, + ], + }, + }); + expect(await findLatestParentMessageId("t-fb")).toBe("msg-Y"); + }); + + it("returns null when state fetch fails", async () => { + mockGetState.mockRejectedValueOnce(new Error("network down")); + expect(await findLatestParentMessageId("t-err")).toBeNull(); }); - it("returns null for an unknown thread", async () => { - expect(await findLatestParentMessageId("ghost-thread")).toBeNull(); + it("returns null when state has no messages", async () => { + mockGetState.mockResolvedValueOnce({ values: {} }); + expect(await findLatestParentMessageId("t-nomsgs")).toBeNull(); }); });