[Feat]: tighten KB Indexed badge + folder cascade delete + pmid attribution

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<number>`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<number>`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,
   });

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,

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 });
   }

backend/node/kb/page-to-markdown-node.ts

@@ -21,7 +21,18 @@ export const ocrPageSchema = z.object({
 export async function pageToMarkdownNode(
   state: KbAgentStateShape,
 ): Promise<Partial<KbAgentStateShape>> {
-  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,

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");
   }

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<Partial<KbAgentStateShape>> {
+  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,
   );

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,
   };
 }
 

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 = {

components/assistant-ui/kb-mention.tsx

@@ -226,15 +226,15 @@ const KbMentionPopoverHeader = () => {
 
   if (scope.activeCategoryId) {
     return (
-      <ComposerPrimitive.Unstable_TriggerPopoverBack className="flex w-full items-center gap-1.5 px-2.5 py-1.5 text-[10px] font-bold capitalize tracking-wider text-muted-foreground/60 border-b border-border/40 mb-1 hover:text-foreground transition-colors cursor-pointer sticky top-0 bg-popover z-10 shrink-0">
+      <ComposerPrimitive.Unstable_TriggerPopoverBack className="flex w-full items-center gap-1.5 px-2.5 py-2 text-[10px] font-bold capitalize tracking-wider text-muted-foreground/60 border-b border-border/40 mb-1 hover:text-foreground transition-colors cursor-pointer sticky top-0 bg-popover z-10 shrink-0">
         <ChevronLeftIcon className="size-3" />
         <span>{"Back"}</span>
       </ComposerPrimitive.Unstable_TriggerPopoverBack>
     );
   }
 
   return (
-    <div className="px-2.5 py-1.5 text-[10px] font-bold capitalize tracking-wider text-muted-foreground/60 border-b border-border/40 mb-1 sticky top-0 bg-popover z-10 shrink-0">
+    <div className="px-2.5 py-2 text-[10px] font-bold capitalize tracking-wider text-muted-foreground/60 border-b border-border/40 mb-1 sticky top-0 bg-popover z-10 shrink-0">
       Knowledge Base
     </div>
   );
@@ -288,7 +288,7 @@ export const KbMentionPopover = memo(function KbMentionPopover() {
       <ComposerPrimitive.Unstable_TriggerPopover
         char="@"
         adapter={bundle.adapter}
-        className="aui-composer-trigger-popover bg-popover text-popover-foreground absolute bottom-full left-0 z-50 mb-4 w-72 max-h-72 overflow-y-auto rounded-xl border bg-background p-1 shadow-lg flex flex-col gap-0.5"
+        className="aui-composer-trigger-popover bg-popover text-popover-foreground absolute bottom-full left-0 z-50 mb-4 w-72 max-h-72 overflow-y-auto rounded-xl border bg-background px-1 shadow-lg flex flex-col gap-0.5 pb-1"
       >
         <PopoverOpenWatcher onOpen={refetch} />
         {/* Directive must always be mounted so the popover can open when @ is typed */}