From 9e39a9c8c634aae5245564ede3ddcccf3aeaba59 Mon Sep 17 00:00:00 2001 From: FireTable Date: Thu, 23 Jul 2026 13:21:20 +0800 Subject: [PATCH 1/8] fix(observability): resolve parent_message_id per-call from metadata, drop shared instance state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CapturingHandler.currentParentMessageId was a single instance field overwritten by every outermost handleChainStart. Concurrent invokes (main chat + a background_agent / kbAgent dispatched mid-turn) clobbered each other, and inner spans inherited whichever message arrived last — kb-upload inner spans got attributed to the user's NEXT message instead of the trigger. Drop the instance field. handleChainStart now resolves parent_message_id per-call: prefer metadata.parent_message_id (per-run, set by runs.create on kbAgent / background_agent dispatches), fallback to lastHumanMessageId(inputs.messages) for the chat streaming path. Inner LLM / tool spans that don't get either source stay null at capture; bulkInsertSpans' state-fallback backfill fills them at INSERT time from langGraphClient.threads.getState. No shared state, no clobber race. --- lib/observability/callback.ts | 65 ++++++----- lib/observability/queries.ts | 61 ++++++---- .../observability/callback-collector.test.ts | 80 ++++++++++++- tests/lib/observability/queries.test.ts | 110 +++++++++++------- 4 files changed, 222 insertions(+), 94 deletions(-) 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..231f2344 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 @@ -124,6 +126,12 @@ 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. + + console.warn( + 2222, + rows.map((item) => ({ threadId: item.threadId, parentMessageId: item.parentMessageId })), + ); + await backfillWaitingInterruptSpans(rows); // ponytail: ON CONFLICT DO NOTHING makes the write idempotent — the // same runId can fire bulkInsertSpans twice (parent chain end + @@ -136,32 +144,45 @@ 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: 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) { + if (typeof r.meta?.parent_message_id === "string") { + r.parentMessageId = r.meta?.parent_message_id as string; + } else { + missingThreads.add(r.threadId); + } + } } + for (const tid of missingThreads) { const latest = await findLatestParentMessageId(tid); if (!latest) continue; 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/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(); }); }); From feb73bb5786a9d132a6cba3a514846ba5206c4b1 Mon Sep 17 00:00:00 2001 From: FireTable Date: Thu, 23 Jul 2026 13:57:45 +0800 Subject: [PATCH 2/8] feat(kb): tighten Indexed badge to require vectors + entity/relationship rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ChunksStatusBadge used to flip to 'Indexed' as soon as every chunk had status='success'. status='success' alone doesn't mean the chunk is queryable: pgvector needs the vector, and the entity-leg of hybrid search needs kb_entity / kb_relationship rows. A doc with all vectors done but no entities yet (extraction still running) used to be 'Indexed' — now it sits in 'Embedding'. Redefine successChunks to mean status='success' AND embedding IS NOT NULL. Add embeddingPendingChunks + entityCount + relationshipCount aggregates (two new subqueries on kb_entity / kb_relationship grouped by document_id). Badge collapses three in-flight lanes into 'Extracting' (chunks being OCR'd / split) and 'Embedding' (chunks done, vector or graph-extract still running); 'Indexed' now requires all three legs of hybrid search populated. --- app/api/kb/documents/[id]/route.ts | 27 ++++- app/api/kb/documents/route.ts | 6 ++ backend/tool/kb/list-documents.ts | 6 ++ backend/tool/kb/types.ts | 3 + .../settings/kb-view/doc-detail-dialog.tsx | 19 +++- components/settings/kb-view/doc-table.tsx | 48 +++++---- components/settings/kb-view/types.ts | 7 ++ components/tool-ui/kb/list-documents-card.tsx | 3 + components/tool-ui/kb/status-badge.tsx | 60 +++++++++-- components/tool-ui/kb/types.ts | 3 + lib/kb/queries.ts | 99 +++++++++++++++++-- lib/kb/source-kind.ts | 2 +- .../frontend/kb-list-documents-card.test.tsx | 3 + 13 files changed, 242 insertions(+), 44 deletions(-) diff --git a/app/api/kb/documents/[id]/route.ts b/app/api/kb/documents/[id]/route.ts index e68b0652..db73a0bd 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, graphCounts] = 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). + db + .select({ + entityCount: sql`count(${kbEntity.id})::int`.as("entity_count"), + relationshipCount: sql`count(${kbRelationship.id})::int`.as("relationship_count"), + }) + .from(kbEntity) + .leftJoin(kbRelationship, eq(kbEntity.documentId, kbRelationship.documentId)) + .where(and(eq(kbEntity.documentId, doc.id), eq(kbEntity.userId, user.id))) + .groupBy(kbEntity.documentId), + ]); + + const entityCount = graphCounts[0]?.entityCount ?? 0; + const relationshipCount = graphCounts[0]?.relationshipCount ?? 0; 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/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/settings/kb-view/doc-detail-dialog.tsx b/components/settings/kb-view/doc-detail-dialog.tsx index a6cf34a6..beeed716 100644 --- a/components/settings/kb-view/doc-detail-dialog.tsx +++ b/components/settings/kb-view/doc-detail-dialog.tsx @@ -277,10 +277,22 @@ export function DocDetailDialog({ if (!d) return null; const totalChunks = d.chunks.length; - const successChunks = d.chunks.filter((c) => c.status === "success").length; + // ponytail: successChunks now means status='success' AND hasEmbedding — + // a chunk with status='success' but no pgvector vector isn't queryable + // by hybrid search, so it shouldn't count toward "Indexed". Those land + // in embeddingPendingChunks instead and show as "Embedding" in the UI. + const successChunks = d.chunks.filter((c) => c.status === "success" && c.hasEmbedding).length; + const embeddingPendingChunks = d.chunks.filter( + (c) => c.status === "success" && !c.hasEmbedding, + ).length; const failedChunks = d.chunks.filter((c) => c.status === "failed").length; const pendingChunks = d.chunks.filter((c) => c.status === "pending").length; const parsingChunks = d.chunks.filter((c) => c.status === "parsing").length; + // ponytail: doc-detail gets these from the API route (parallel count + // query alongside the chunks fetch); the badge requires the third leg + // of hybrid search to be populated before showing "Indexed". + const entityCount = d.doc.entityCount ?? 0; + const relationshipCount = d.doc.relationshipCount ?? 0; // ponytail: derive page-level counts the same way chunksStatusBadge // does. When pages carry an explicit `status` mirror it; legacy rows @@ -352,9 +364,12 @@ export function DocDetailDialog({ @@ -363,7 +378,7 @@ export function DocDetailDialog({ variant="outline" className="border-none bg-transparent text-muted-foreground shadow-none px-0 py-0.5 font-normal leading-none" > - {mimeShortLabel(d.doc.contentType)} + {mimeShortLabel(d.doc.contentType)} {d.doc.pages && diff --git a/components/settings/kb-view/doc-table.tsx b/components/settings/kb-view/doc-table.tsx index e5a5f6d0..f34f52ea 100644 --- a/components/settings/kb-view/doc-table.tsx +++ b/components/settings/kb-view/doc-table.tsx @@ -253,9 +253,12 @@ export function DocRow({ @@ -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/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..445b73ff 100644 --- a/components/tool-ui/kb/status-badge.tsx +++ b/components/tool-ui/kb/status-badge.tsx @@ -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; }) { @@ -171,10 +186,21 @@ export function ChunksStatusBadge({ ); } else { - const terminal = success + failed; - const isCompleted = success === total; + const embeddingPending = embeddingPendingChunks ?? 0; + const entities = entityCount ?? 0; + const relationships = relationshipCount ?? 0; + const hasGraph = entities > 0 || relationships > 0; + // ponytail: terminal = chunks that finished OCR (success + embedPending) + // or hit a terminal failure. "Indexed" needs ALL THREE pipeline + // outputs ready: vector + entity/relationship. If chunks all have + // vectors but no entities/relationships yet, that's still "Embedding" + // — the entity-leg of hybrid search has nothing to score against. + const terminal = success + failed + embeddingPending; + const isCompleted = success === total && hasGraph; const isFailed = failed > 0 && terminal === total; - const isIndexing = !isCompleted && !isFailed; + const hasInflightChunks = pending > 0 || parsing > 0; + const hasInflightExtraction = !hasGraph; // vectors done but no graph rows yet + const isInProgress = !isCompleted && !isFailed; let variant: "success" | "destructive" | "muted" = "muted"; if (isFailed) { @@ -186,9 +212,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; 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/lib/kb/queries.ts b/lib/kb/queries.ts index 3bd27960..647252fc 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) ?? [], @@ -1442,12 +1506,31 @@ 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 }, + }); + + for (const doc of docs) { + const threadId = doc.id.replace(/^d-/, ""); + await tx.delete(threads).where(and(eq(threads.id, threadId), eq(threads.kind, "kb"))); + } + + 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..52ed9f2a 100644 --- a/lib/kb/source-kind.ts +++ b/lib/kb/source-kind.ts @@ -70,7 +70,7 @@ 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); + return TYPE_LABEL[subtype] ?? mimeType; } // ponytail: turn `R2_ALLOWED_CONTENT_TYPES` (comma-separated mimes) 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, }; } From a73a179c015698a16b64b5ff43072eaf45a9d60e Mon Sep 17 00:00:00 2001 From: FireTable Date: Thu, 23 Jul 2026 14:12:40 +0800 Subject: [PATCH 3/8] feat(kb): reprocess skip + folder cascade delete + dialog UX MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - skipFileToPageNode + pageToMarkdownNode also skip on retryFailed / retryFailedChunks (was chunksOnly only); doc-status preservation test - prepareKBDataNode reads config.configurable.parent_message_id first with lastHumanMessageId fallback - DELETE /api/kb/folders/[id] accepts ?deleteAll=true to cascade-delete every doc in the folder (each doc's chunks/embeddings/attachments) before the folder row; Settings dialog adds a switch + typed "delete all" confirmation - DocDeleteDialog / DocReprocessDialog close themselves via onOpenChange before firing onDeleted / onReprocessed; long titles use break-all - ObservabilityPopover no longer auto-closes on row click — the Sheet is the singleton entry, not the popover - kb-mention popover tweaks: scrollable body, sticky header, padding - folder-sidebar callbacks no longer reset selection on submit (the dialog does it via onOpenChange) - docs: APIS.md documents the new indexing-count fields on GET /api/kb/documents + GET /api/kb/documents/[id], the chunks[].hasEmbedding flag, and the folder ?deleteAll cascade. OBSERVABILITY.md documents per-call pmid resolution from metadata.parent_message_id + langGraphClient.threads.getState fallback. TOOLS.md documents the expanded list_documents payload shape shared with the Settings badges. --- app/api/kb/folders/[id]/route.ts | 17 ++- backend/node/kb/page-to-markdown-node.ts | 6 +- backend/node/kb/prepare-kb-data-node.ts | 3 +- backend/node/kb/split-file-page-node.ts | 11 ++ components/assistant-ui/kb-mention.tsx | 6 +- components/settings/kb-view/dialogs.tsx | 135 +++++++++++++++--- .../settings/kb-view/folder-sidebar.tsx | 10 +- .../kb-view/observability-popover.tsx | 1 - docs/APIS.md | 37 +++-- docs/OBSERVABILITY.md | 4 +- docs/TOOLS.md | 14 +- tests/api/kb/folder.test.ts | 40 +++++- tests/api/kb/reprocess.test.ts | 46 ++++++ 13 files changed, 275 insertions(+), 55 deletions(-) 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..22c28e29 100644 --- a/backend/node/kb/page-to-markdown-node.ts +++ b/backend/node/kb/page-to-markdown-node.ts @@ -21,7 +21,11 @@ export const ocrPageSchema = z.object({ export async function pageToMarkdownNode( state: KbAgentStateShape, ): Promise> { - if (state.mode === "chunksOnly") { + if ( + state.mode === "chunksOnly" || + state.mode === "retryFailed" || + state.mode === "retryFailedChunks" + ) { 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/components/assistant-ui/kb-mention.tsx b/components/assistant-ui/kb-mention.tsx index 4c96923d..694515c2 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..2c327632 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,154 @@ 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 +

    +
    + { + setDeleteAll(c); + if (!c) setConfirmText(""); + }} + /> +
    + + {deleteAll && ( +
    + + setConfirmText(e.target.value)} + placeholder="delete all" + className="h-8 font-mono text-xs" + /> +
    + )} +
    + +
0 && terminal === total; + // 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; diff --git a/lib/kb/queries.ts b/lib/kb/queries.ts index a05c8a84..3bac7a7b 100644 --- a/lib/kb/queries.ts +++ b/lib/kb/queries.ts @@ -1527,6 +1527,12 @@ export async function deleteKbFolderForUser( 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. diff --git a/lib/observability/queries.ts b/lib/observability/queries.ts index b5bbdf12..7440a371 100644 --- a/lib/observability/queries.ts +++ b/lib/observability/queries.ts @@ -164,6 +164,18 @@ export async function findLatestParentMessageId(threadId: string): Promise(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, ask // findLatestParentMessageId (which now reads thread state via // langGraphClient) for the latest HumanMessage id on that thread and @@ -173,21 +185,26 @@ async function backfillParentMessageIds(rows: NewObservabilitySpanRow[]): Promis const missingThreads = new Set(); for (const r of rows) { if (r.parentMessageId == null) { - if (typeof r.meta?.parent_message_id === "string") { - r.parentMessageId = r.meta?.parent_message_id as string; - } else { - missingThreads.add(r.threadId); - } + 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/backend/node/kb/page-to-markdown-node.test.ts b/tests/backend/node/kb/page-to-markdown-node.test.ts index c13bba97..a965783b 100644 --- a/tests/backend/node/kb/page-to-markdown-node.test.ts +++ b/tests/backend/node/kb/page-to-markdown-node.test.ts @@ -89,4 +89,51 @@ describe("pageToMarkdownNode — retryFailed mode does NOT skip OCR", () => { 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"); + }); }); From d70418b7ec643005b0f912f208584af8377c56f8 Mon Sep 17 00:00:00 2001 From: FireTable Date: Thu, 23 Jul 2026 15:14:13 +0800 Subject: [PATCH 8/8] test(observability): mock langGraphClient.threads.getState in CI tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - bulkInsertSpans → backfillParentMessageIds → findLatestParentMessageId hits langGraphClient.threads.getState. CI has no langgraph dev server, so the call hangs and 9 bulkInsertSpans tests time out at 5s. Stub the HTTP layer in the two affected test files so the state fetch returns immediately (no human messages → null pmid → row persists pmid-less, the panel 404s those by design). - The 500ms timeout in backfillParentMessageIds is kept as defense-in-depth for the real prod server hanging — but the test env no longer depends on it. --- tests/api/threads/observability.test.ts | 20 ++++++++++++++++++- .../integration/observability-wiring.test.ts | 16 ++++++++++++++- 2 files changed, 34 insertions(+), 2 deletions(-) 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/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";