From 42a33790ae742902e6ecaef7e716f15bd465cbc2 Mon Sep 17 00:00:00 2001 From: mimran-khan Date: Thu, 12 Nov 2026 02:50:50 +0530 Subject: [PATCH 1/2] feat: open document citations in the page viewer A citation in an answer opens the references viewer on the cited page. The references rail lists those pages, and a citation the model split across two lines stays one inline link. Signed-off-by: mimran-khan --- src/frontend/components/ChatMessagesView.tsx | 74 ++- .../components/DocumentReferences.tsx | 117 +++++ .../components/SourceDocumentViewer.tsx | 265 +++++++++++ .../components/document-open-context.ts | 3 + .../components/document-open.test.tsx | 150 ++++++ src/frontend/global.css | 13 +- src/frontend/hooks/useStreamingAPI.ts | 23 + src/frontend/lib/document-references.test.ts | 153 ++++++ src/frontend/lib/document-references.ts | 448 ++++++++++++++++++ src/frontend/lib/processed-document.test.ts | 89 ++++ src/frontend/lib/processed-document.ts | 216 +++++++++ src/server/citation-images.test.ts | 28 ++ src/server/citation-images.ts | 33 ++ 13 files changed, 1591 insertions(+), 21 deletions(-) create mode 100644 src/frontend/components/DocumentReferences.tsx create mode 100644 src/frontend/components/SourceDocumentViewer.tsx create mode 100644 src/frontend/components/document-open-context.ts create mode 100644 src/frontend/components/document-open.test.tsx create mode 100644 src/frontend/lib/document-references.test.ts create mode 100644 src/frontend/lib/document-references.ts create mode 100644 src/frontend/lib/processed-document.test.ts create mode 100644 src/frontend/lib/processed-document.ts create mode 100644 src/server/citation-images.test.ts create mode 100644 src/server/citation-images.ts diff --git a/src/frontend/components/ChatMessagesView.tsx b/src/frontend/components/ChatMessagesView.tsx index e39063a2..daa61755 100644 --- a/src/frontend/components/ChatMessagesView.tsx +++ b/src/frontend/components/ChatMessagesView.tsx @@ -4,13 +4,12 @@ import { AlertCircle, Check, ChevronDown, ChevronRight, Copy, Download, Loader2, import type { InterruptInfo } from "../types/deep-agent"; import { InputForm } from "./InputForm"; import { McpStatusPanel } from "./McpStatusPanel"; -import { useState, ReactNode, useMemo, useEffect, useRef, useCallback, cloneElement, isValidElement } from "react"; +import { useState, ReactNode, useMemo, useEffect, useRef, useCallback, useContext, cloneElement, isValidElement } from "react"; import { cn } from "../lib/utils"; import { Dropdown, DropdownItem, DropdownList, - Label, MenuToggle, type MenuToggleElement, } from "@patternfly/react-core"; @@ -26,6 +25,9 @@ import { markdownSanitizeSchema } from "../lib/markdown-sanitize"; import { isSubAgentToolCall, type SubagentTranscriptEntry } from "../types/deep-agent"; import { SubAgentIndicator } from "./SubAgentIndicator"; import { ToolCallRow } from "./ToolCallRow"; +import { DocumentOpenContext } from "./document-open-context"; +import { DocumentChrome } from "./DocumentReferences"; +import { fileIdFromCitationHref, isInAppHref, joinSplitCitationLinks } from "../lib/document-references"; import { ChartCard, PlotlyCard, useContentSegments } from "./InlineChart"; import { TodoStrip } from "./TodoStrip"; import { FeedbackButtons, FeedbackCommentBox, useFeedback } from "./FeedbackButtons"; @@ -287,6 +289,37 @@ function MessageActionsRow({ ); } +function CitationLink({ className, children, href }: MdComponentProps) { + const openDocument = useContext(DocumentOpenContext); + const link = typeof href === "string" ? href : ""; + const fileId = fileIdFromCitationHref(link); + if (fileId) { + return ( + + ); + } + if (isInAppHref(link, globalThis.location?.origin ?? "")) { + return {children}; + } + return ( + + {children} + (opens in new tab) + + ); +} + const mdComponents = { h1: ({ className, children, ...props }: MdComponentProps) => (

@@ -308,20 +341,7 @@ const mdComponents = { {children}

), - a: ({ className, children, href, ...props }: MdComponentProps) => ( - - ), + a: CitationLink, ul: ({ className, children, ...props }: MdComponentProps) => (
    {children} @@ -760,7 +780,7 @@ interface AIMessageRendererProps { } function AnswerStream({ text }: { text: string }) { - const parts = useContentSegments(text); + const parts = useContentSegments(joinSplitCitationLinks(text)); return (
    {parts.map((part, index) => { @@ -778,6 +798,20 @@ function AnswerStream({ text }: { text: string }) { export function AIMessageRenderer({ message, pendingInterrupt, onInterruptResume, onAlwaysAllow, globalApprovalIndex = 0, approvalSlotOffset = 0, totalActionRequests = 0, allDecisionsMade = false, onSingleDecision, thinkingLive = false, questionContent, interruptContent, turnMessages }: AIMessageRendererProps) { const source = turnMessages && turnMessages.length > 0 ? turnMessages : [message]; + const openDocument = useContext(DocumentOpenContext); + const keepDocumentInPage = useCallback((event: React.MouseEvent) => { + const target = event.target; + if (!(target instanceof Element)) return; + const anchor = target.closest("a"); + if (!(anchor instanceof HTMLAnchorElement)) return; + const href = anchor.getAttribute("href") ?? ""; + const fileId = fileIdFromCitationHref(href); + const staysInChat = fileId !== null || isInAppHref(href, globalThis.location?.origin ?? ""); + if (!staysInChat) return; + event.preventDefault(); + event.stopPropagation(); + if (fileId) openDocument?.(fileId, href); + }, [openDocument]); const [thoughtChoice, setThoughtChoice] = useState>({}); const messageKey = JSON.stringify(source); @@ -932,7 +966,7 @@ export function AIMessageRenderer({ message, pendingInterrupt, onInterruptResume if (item.kind === 'thinking') { return (
    - {item.text} + {joinSplitCitationLinks(item.text)}
    ); } @@ -999,7 +1033,7 @@ export function AIMessageRenderer({ message, pendingInterrupt, onInterruptResume }, [messageKey, thoughtChoice, globalApprovalIndex, allDecisionsMade, pendingInterrupt, onInterruptResume, onAlwaysAllow, onSingleDecision, totalActionRequests, approvalSlotByToolCallId, thinkingLive, questionContent, interruptContent, pendingToolNames]); return ( -
    +
    {renderMessage}
    ); @@ -1193,6 +1227,7 @@ export function ChatMessagesView({ const showExport = Boolean(onExportMarkdown && onExportJson && messages.length > 0); return ( +
    +     ); } diff --git a/src/frontend/components/DocumentReferences.tsx b/src/frontend/components/DocumentReferences.tsx new file mode 100644 index 00000000..d92c9152 --- /dev/null +++ b/src/frontend/components/DocumentReferences.tsx @@ -0,0 +1,117 @@ +import { useCallback, useEffect, useMemo, useState, type ReactNode } from "react"; +import { BookOpen, ChevronRight, FileText } from "lucide-react"; +import { documentReferences, viewerPages, type DocumentReference } from "../lib/document-references"; +import { prefetchCitedPages } from "../lib/processed-document"; +import { callMcpAppTool } from "../services/mcp-apps-api"; +import { DocumentOpenContext } from "./document-open-context"; +import { SourceDocumentViewer } from "./SourceDocumentViewer"; + +function pageLabel(pages: number[]): string { + if (pages.length === 0) return ""; + return `p. ${pages.join(", ")}`; +} + +export function ReferencesRail({ + references, + activeFileId, + onOpen, +}: { + references: DocumentReference[]; + activeFileId: string | null; + onOpen: (fileId: string) => void; +}) { + const [expanded, setExpanded] = useState(true); + if (references.length === 0) return null; + if (!expanded) { + return ( + + ); + } + return ( + + ); +} + +export function DocumentChrome({ + messages, + children, +}: { + messages: readonly unknown[]; + children: ReactNode; +}) { + const references = useMemo(() => documentReferences(messages), [messages]); + const pages = useMemo(() => viewerPages(references), [references]); + useEffect(() => { + prefetchCitedPages(pages, callMcpAppTool); + }, [pages]); + const [openDoc, setOpenDoc] = useState<{ fileId: string; href?: string } | null>(null); + const [activeFileId, setActiveFileId] = useState(null); + const openDocument = useCallback((fileId: string, href?: string) => { + setOpenDoc({ fileId, href }); + setActiveFileId(fileId); + }, []); + const showFile = useCallback((fileId: string) => setActiveFileId(fileId), []); + return ( + +
    +
    {children}
    + +
    + {openDoc ? ( + { + setOpenDoc(null); + setActiveFileId(null); + }} + onActiveFile={showFile} + onOpenDocument={openDocument} + /> + ) : null} +     + ); +} diff --git a/src/frontend/components/SourceDocumentViewer.tsx b/src/frontend/components/SourceDocumentViewer.tsx new file mode 100644 index 00000000..b2770433 --- /dev/null +++ b/src/frontend/components/SourceDocumentViewer.tsx @@ -0,0 +1,265 @@ +import { useEffect, useMemo, useRef, useState } from "react"; +import { createPortal } from "react-dom"; +import ReactMarkdown from "react-markdown"; +import remarkGfm from "remark-gfm"; +import rehypeSanitize from "rehype-sanitize"; +import { ChevronLeft, ChevronRight, FileText, RotateCcw, X, ZoomIn, ZoomOut } from "lucide-react"; +import { markdownSanitizeSchema } from "../lib/markdown-sanitize"; +import { fileIdFromCitationHref, isInAppHref, viewerPages, viewerStartIndex, type DocumentReference } from "../lib/document-references"; +import { loadPageImages, peekPageImage } from "../lib/processed-document"; +import { callMcpAppTool } from "../services/mcp-apps-api"; +import type { CitationImage } from "../../server/citation-images"; + +function excerpt(text: string): string { + const plain = text.replace(/^#+\s*/gm, "").replace(/\s+/g, " ").trim(); + if (plain.length <= 140) return plain; + return `${plain.slice(0, 137)}...`; +} + +function ViewerToolbar({ + title, + pageLabel, + zoom, + onZoomIn, + onZoomOut, + onReset, + onPrev, + onNext, + canPrev, + canNext, + showNav, + onClose, +}: { + title: string; + pageLabel: string; + zoom: number; + onZoomIn: () => void; + onZoomOut: () => void; + onReset: () => void; + onPrev: () => void; + onNext: () => void; + canPrev: boolean; + canNext: boolean; + showNav: boolean; + onClose: () => void; +}) { + return ( +
    +
    +
    +
    + {showNav ? ( + <> + + +
    + + ) : null} + + {Math.round(zoom * 100)}% + + +
    + +
    +
    + ); +} + +function PageImage({ title, image }: { title: string; image: CitationImage }) { + return ( + {`${title} + ); +} + +function DocumentPaper({ + title, + markdown, + image, + loading, + onOpenDocument, +}: { + title: string; + markdown: string; + image: CitationImage | null; + loading: boolean; + onOpenDocument?: (fileId: string, href?: string) => void; +}) { + if (image) return ; + if (markdown) { + return ; + } + if (loading) { + return

    Loading the page...

    ; + } + return

    This document was not in the retrieved passages.

    ; +} + +function DocumentArticle({ + markdown, + onOpenDocument, +}: { + markdown: string; + onOpenDocument?: (fileId: string, href?: string) => void; +}) { + return ( +
    + {markdown ? ( + { + const link = typeof href === "string" ? href : ""; + const fileId = fileIdFromCitationHref(link); + if (fileId && onOpenDocument) { + return ( + + ); + } + if (fileId || isInAppHref(link, globalThis.location?.origin ?? "")) { + return {children}; + } + return {children}; + }, + }} + > + {markdown} + + ) : null} +
    + ); +} + +function pageCaption(page: number | null, index: number, count: number): string { + const where = count > 1 ? `${index + 1} of ${count}` : ""; + const label = page != null ? `Page ${page}` : "Document"; + return [label, where].filter(Boolean).join(" · "); +} + +export function SourceDocumentViewer({ + fileId, + href = "", + references, + onClose, + onActiveFile, + onOpenDocument, +}: { + fileId: string; + href?: string; + references: readonly DocumentReference[]; + onClose: () => void; + onActiveFile?: (fileId: string) => void; + onOpenDocument?: (fileId: string, href?: string) => void; +}) { + const pages = useMemo(() => viewerPages(references), [references]); + const pagesRef = useRef(pages); + pagesRef.current = pages; + const [index, setIndex] = useState(() => viewerStartIndex(pages, fileId, href)); + const [zoom, setZoom] = useState(1); + const [readyTick, setReadyTick] = useState(0); + const current = pages[index] ?? pages[0]; + const title = current?.title ?? "Document"; + const markdown = current?.text ?? ""; + const storedImage = current?.pipeline && current.chunkId + ? peekPageImage(current.pipeline, current.chunkId) + : null; + const image = storedImage ?? null; + const waitingForImage = Boolean(current?.pipeline && current.chunkId && storedImage === undefined && readyTick >= 0); + + useEffect(() => { + setIndex(viewerStartIndex(pagesRef.current, fileId, href)); + }, [fileId, href]); + + useEffect(() => { + const onKey = (event: KeyboardEvent) => { + if (event.key === "Escape") onClose(); + if (event.key === "ArrowRight") setIndex((value) => Math.min(value + 1, Math.max(pagesRef.current.length - 1, 0))); + if (event.key === "ArrowLeft") setIndex((value) => Math.max(value - 1, 0)); + if (event.key === "+" || event.key === "=") setZoom((value) => Math.min(value + 0.25, 3)); + if (event.key === "-") setZoom((value) => Math.max(value - 0.25, 0.5)); + }; + window.addEventListener("keydown", onKey); + return () => window.removeEventListener("keydown", onKey); + }, [onClose]); + + useEffect(() => { + const page = pages[index]; + if (!page) return; + onActiveFile?.(page.fileId); + if (!page.pipeline || !page.chunkId) return; + if (peekPageImage(page.pipeline, page.chunkId) !== undefined) return; + let active = true; + loadPageImages(page.pipeline, page.chunkId, [], callMcpAppTool).finally(() => { + if (active) setReadyTick((value) => value + 1); + }); + return () => { + active = false; + }; + }, [index, pages, onActiveFile]); + + const cited = current?.text ? excerpt(current.text) : ""; + const pageLabel = pageCaption(current?.page ?? null, index, pages.length); + + return createPortal( +
    + setZoom((value) => Math.min(value + 0.25, 3))} + onZoomOut={() => setZoom((value) => Math.max(value - 0.25, 0.5))} + onReset={() => setZoom(1)} + onPrev={() => setIndex((value) => Math.max(value - 1, 0))} + onNext={() => setIndex((value) => Math.min(value + 1, Math.max(pages.length - 1, 0)))} + canPrev={index > 0} + canNext={index < pages.length - 1} + showNav={pages.length > 1} + onClose={onClose} + /> +
    +
    + +
    +
    + {cited ? ( +
    +
    +

    + Cited passage + {": "} + {cited} +

    +
    + ) : null} +
    , + document.body, + ); +} diff --git a/src/frontend/components/document-open-context.ts b/src/frontend/components/document-open-context.ts new file mode 100644 index 00000000..3e623388 --- /dev/null +++ b/src/frontend/components/document-open-context.ts @@ -0,0 +1,3 @@ +import { createContext } from "react"; + +export const DocumentOpenContext = createContext<((fileId: string, href?: string) => void) | null>(null); diff --git a/src/frontend/components/document-open.test.tsx b/src/frontend/components/document-open.test.tsx new file mode 100644 index 00000000..5b7ec494 --- /dev/null +++ b/src/frontend/components/document-open.test.tsx @@ -0,0 +1,150 @@ +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import type { Message } from "@langchain/langgraph-sdk"; +import { callMcpAppTool } from "../services/mcp-apps-api"; +import { AIMessageRenderer } from "./ChatMessagesView"; +import { DocumentChrome } from "./DocumentReferences"; + +vi.mock("../services/mcp-apps-api", () => ({ + callMcpAppTool: vi.fn(), +})); + +const doc = "https://docs.google.com/document/d/doc-1"; + +describe("document open", () => { + it("opens the document in the page viewer", async () => { + vi.mocked(callMcpAppTool).mockImplementation(async (_mcp, name) => { + if (name === "get_processed_document") { + return { + content: [{ + type: "text", + text: `Processed document:\n${JSON.stringify({ markdown_content: "## Principles\n\nStandardize on RHEL." })}`, + }], + }; + } + return { content: [] }; + }); + const user = userEvent.setup(); + const message = { + type: "ai", + id: "m1", + content: `See [CY26 RHEL](${doc}).`, + tool_calls: [{ + name: "get_chunks_for_embeddings", + content: `Found 1 chunks for query in pipeline "analyzeproductrhel":\n${JSON.stringify([{ + file_id: "doc-1", + chunk_text: "## Principles\nStandardize on RHEL.", + }])}\n`, + }], + } as Message; + render( + + + , + ); + expect(screen.queryAllByRole("link")).toHaveLength(0); + const rail = screen.getByRole("complementary", { name: "References" }); + await user.click(within(rail).getByRole("button", { name: /CY26 RHEL/ })); + expect(document.querySelector("iframe")).toBeNull(); + expect(await screen.findByRole("dialog", { name: "CY26 RHEL" })).toBeInTheDocument(); + expect(await screen.findByRole("heading", { name: "Principles" })).toBeInTheDocument(); + expect(screen.getByText("Standardize on RHEL.")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Close document" })); + expect(screen.queryByRole("dialog")).toBeNull(); + }); + + it("opens the viewer from a bare file id and does not follow a relative reference", async () => { + const fileId = "1kay0XsAOelA-mJ1kIWvXC2KAAKonMurqxMNIZ0yL95s"; + const user = userEvent.setup(); + const message = { + type: "ai", + id: "m-bare", + content: `See [CY26 RHEL](${fileId}). The [\`flag\` reference](./calculate_bookings_flag.md) stays here.`, + tool_calls: [{ + name: "get_chunks_for_embeddings", + content: `Found 1 chunks for query in pipeline "analyzeproductrhel":\n${JSON.stringify([{ + file_id: fileId, + chunk_text: "Standardize on RHEL.", + }])}\n`, + }], + } as Message; + render( + + + , + ); + expect(screen.queryByRole("link", { name: /reference/ })).toBeNull(); + expect(screen.getByText(/reference/)).toBeInTheDocument(); + const citations = screen.getAllByRole("button", { name: "CY26 RHEL" }); + expect(citations[0]).toHaveClass("citation-link"); + await user.click(citations[0]); + expect(await screen.findByRole("dialog", { name: "CY26 RHEL" })).toBeInTheDocument(); + }); + + it("renders the citation page from the returned image", async () => { + vi.mocked(callMcpAppTool).mockImplementation(async (_mcp, name) => { + if (name === "get_citation_image") { + return { content: [{ type: "image", data: "abc", mimeType: "image/png" }] }; + } + return { content: [] }; + }); + const user = userEvent.setup(); + const message = { + type: "ai", + id: "m2", + content: `See [CY26 RHEL](${doc}).`, + tool_calls: [{ + name: "get_chunks_for_embeddings", + content: `Found 1 chunks for query in pipeline "productstrategycy26":\n${JSON.stringify([{ + file_id: "doc-1", + chunk_id: "pipelines/doc.pdf#1", + chunk_text: "page text", + regions: [{ pageNo: 1, bbox: [0, 0, 1, 1], image: "page.png" }], + }])}\n`, + }], + } as Message; + render( + +
    + , + ); + await user.click(within(screen.getByRole("complementary", { name: "References" })).getByRole("button", { name: /CY26 RHEL/ })); + const image = await screen.findByRole("img", { name: "CY26 RHEL page" }); + expect(image).toHaveAttribute("src", "data:image/png;base64,abc"); + expect(screen.queryByText("page text")).toBeNull(); + }); + + it("opens a pipeline citation and moves to the next referenced page", async () => { + vi.mocked(callMcpAppTool).mockResolvedValue({ content: [] }); + const edge = "1cNsyB6IoH4rpHM6lKeRdGkgO0Vzvy1f4-RvkRJ_6NmA"; + const rhel = "17dPGPaiFp-eBYOZVtosqtDa03opeliamHhafZVrRVxU"; + const edgeChunk = `pipelines/productstrategycy26/stages/crawl/${edge}.pdf#13`; + const rhelChunk = `pipelines/productstrategycy27/stages/crawl/${rhel}.pdf#28`; + const user = userEvent.setup(); + const message = { + type: "ai", + id: "m-pages", + content: `See [Edge Strategy, p.9](${edgeChunk}) and [RHEL Strategy, p.19](${rhelChunk}).`, + tool_calls: [{ + name: "get_chunks_for_embeddings", + content: `Found 2 chunks for query in pipeline "productstrategycy26":\n${JSON.stringify([ + { file_id: edge, chunk_id: edgeChunk, chunk_text: "Edge page nine", regions: [{ pageNo: 9 }] }, + { file_id: rhel, chunk_id: rhelChunk, chunk_text: "RHEL page nineteen", regions: [{ pageNo: 19 }] }, + ])}\n`, + }], + } as Message; + render( + + + , + ); + await user.click(screen.getByRole("button", { name: "Edge Strategy, p.9" })); + expect(await screen.findByRole("dialog", { name: "Edge Strategy, p.9" })).toBeInTheDocument(); + expect(screen.getByText("Edge page nine")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Next page" })); + expect(screen.getByRole("dialog", { name: "RHEL Strategy, p.19" })).toBeInTheDocument(); + expect(screen.getByText("RHEL page nineteen")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Previous page" })); + expect(screen.getByText("Edge page nine")).toBeInTheDocument(); + }); +}); diff --git a/src/frontend/global.css b/src/frontend/global.css index dddd0633..4ee35f9d 100644 --- a/src/frontend/global.css +++ b/src/frontend/global.css @@ -432,17 +432,26 @@ input[type="url"]:focus { font-style: italic; } -.ai-response-body a { +.ai-response-body a, +button.citation-link { color: var(--primary); text-decoration: underline; text-underline-offset: 2px; text-decoration-color: color-mix(in srgb, var(--primary) 30%, transparent); transition: color 0.15s ease, text-decoration-color 0.15s ease; } -.ai-response-body a:hover { +.ai-response-body a:hover, +button.citation-link:hover { color: color-mix(in srgb, var(--primary) 80%, var(--foreground)); text-decoration-color: color-mix(in srgb, var(--primary) 60%, transparent); } +button.citation-link { + background: none; + border: 0; + padding: 0; + font: inherit; + cursor: pointer; +} .ai-response-body ul { list-style-type: disc; diff --git a/src/frontend/hooks/useStreamingAPI.ts b/src/frontend/hooks/useStreamingAPI.ts index a38d04f7..e2734cc4 100644 --- a/src/frontend/hooks/useStreamingAPI.ts +++ b/src/frontend/hooks/useStreamingAPI.ts @@ -8,7 +8,10 @@ import { type StreamCallback, type StreamStatus, } from '@/lib/streaming/StreamingManager'; +import { citationTargetsInChunkResult, citationTargetsInText } from '@/lib/document-references'; +import { prefetchCitedPages } from '@/lib/processed-document'; import { getStreamingManager } from '@/lib/streaming/streamingManagerRegistry'; +import { callMcpAppTool } from '@/services/mcp-apps-api'; import { useAppDispatch, useAppSelector } from '@/redux/hooks'; import type { AppDispatch } from '@/redux/store'; import { @@ -37,6 +40,18 @@ import { isSubAgentToolCall, extractSubAgentName } from '@/types/deep-agent'; import type { HITLInterruptValue, InterruptInfo } from '@/types/deep-agent'; import { mergeMessageWithMcpModelContext } from '@/types/mcp-apps'; +function warmAnswerText(buffer: { current: string }, content: string, part: 'thinking' | 'text' | undefined): void { + if (part === 'thinking' || !content) return; + buffer.current += content; + const targets = citationTargetsInText(buffer.current); + if (targets.length > 0) prefetchCitedPages(targets, callMcpAppTool); +} + +function warmToolMessage(message: { name?: string; content?: unknown }): void { + const targets = citationTargetsInChunkResult(message.name, message.content); + if (targets.length > 0) prefetchCitedPages(targets, callMcpAppTool); +} + type TokenPart = 'thinking' | 'text'; function dispatchSubagentMessage( @@ -469,6 +484,7 @@ export function useStreamingAPI(threadId: string) { }>({ streamStartTime: null, firstTokenTime: null, streamEndTime: null }); const isStreamingTokensRef = useRef(false); const userCancelledRef = useRef(false); + const citedAnswerRef = useRef(''); const submitLockRef = useRef(false); const streamEndedWithInterruptRef = useRef(false); const pendingInterruptRef = useRef(streamingState.pendingInterrupt); @@ -577,6 +593,7 @@ export function useStreamingAPI(threadId: string) { try { userCancelledRef.current = false; + citedAnswerRef.current = ''; setRetryCount(0); setWasInterrupted(false); const initialMessageCount = messagesRef.current.length; @@ -718,6 +735,7 @@ export function useStreamingAPI(threadId: string) { if (streamClockRef.current.firstTokenTime == null) { streamClockRef.current.firstTokenTime = Date.now(); } + warmAnswerText(citedAnswerRef, content, part); dispatchStreamToken(dispatch, threadId, content, messageId, part, isStreamingTokensRef, subagentToolCallId); }, onDraftDiscard() { @@ -764,6 +782,7 @@ export function useStreamingAPI(threadId: string) { } if (m.type === 'tool') { + warmToolMessage(m); dispatch( mergeToolResult({ chatId: threadId, @@ -1309,6 +1328,7 @@ export function useStreamingAPI(threadId: string) { const callbacks: StreamCallback = { onToken(content, messageId, part, subagentToolCallId) { lastTokenTimeRef.current = Date.now(); + warmAnswerText(citedAnswerRef, content, part); dispatchStreamToken(dispatch, threadId, content, messageId, part, isStreamingTokensRef, subagentToolCallId); }, onDraftDiscard() { @@ -1330,6 +1350,7 @@ export function useStreamingAPI(threadId: string) { return; } if (m.type === 'tool') { + warmToolMessage(m); dispatch( mergeToolResult({ chatId: threadId, @@ -1448,6 +1469,7 @@ export function useStreamingAPI(threadId: string) { onToken(content, messageId, part, subagentToolCallId) { resumeStreamSawProgress = true; lastTokenTimeRef.current = Date.now(); + warmAnswerText(citedAnswerRef, content, part); dispatchStreamToken(dispatch, threadId, content, messageId, part, isStreamingTokensRef, subagentToolCallId); }, onDraftDiscard() { @@ -1468,6 +1490,7 @@ export function useStreamingAPI(threadId: string) { resumeStreamSawProgress = true; if (m.type === 'tool') { + warmToolMessage(m); dispatch( mergeToolResult({ chatId: threadId, diff --git a/src/frontend/lib/document-references.test.ts b/src/frontend/lib/document-references.test.ts new file mode 100644 index 00000000..964d18c8 --- /dev/null +++ b/src/frontend/lib/document-references.test.ts @@ -0,0 +1,153 @@ +import { describe, expect, it } from "vitest"; +import { citationTargetsInChunkResult, citationTargetsInText, documentReferences, fileIdFromCitationHref, fileIdFromDocumentUrl, isInAppHref, joinSplitCitationLinks, viewerPages } from "./document-references"; + +const CHUNK = { + file_id: "doc-1", + chunk_id: "pipelines/a.pdf#4", + chunk_text: "RHEL strategy\nStandardize on security.", + regions: [{ pageNo: 3, bbox: [0, 0, 1, 1], image: "page.png" }], +}; + +describe("joinSplitCitationLinks", () => { + it("puts a wrapped title and file path back on one line", () => { + const href = "pipelines/productstrategycy26/stages/crawl/1nv5kjjHtWqhp6F_QwSuYH3TWdmvoLF0YMZSm7_9oWdM.pdf"; + const raw = `in 2028 [Red Hat AI CY2026 Product Strategy]\n(${href}) .`; + expect(joinSplitCitationLinks(raw)).toBe( + `in 2028 [Red Hat AI CY2026 Product Strategy](${href}) .`, + ); + const wrapped = `[Red\nHat AI CY2026 Product Strategy, p.12]\n(${href})`; + expect(joinSplitCitationLinks(wrapped)).toBe( + `[Red Hat AI CY2026 Product Strategy, p.12](${href})`, + ); + const intact = `See [Edge](${href}#13).`; + expect(joinSplitCitationLinks(intact)).toBe(intact); + }); +}); + +describe("citation warmup", () => { + it("starts from a citation path as soon as the answer contains it", () => { + const fileId = "1cNsyB6IoH4rpHM6lKeRdGkgO0Vzvy1f4-RvkRJ_6NmA"; + const chunkId = `pipelines/productstrategycy26/stages/crawl/${fileId}.pdf#13`; + const partial = `See [Edge Strategy, p.9](${chunkId}`; + expect(citationTargetsInText(partial)).toEqual([ + { pipeline: "productstrategycy26", chunkId }, + ]); + expect(citationTargetsInText("See [flag](./calculate_bookings_flag.md)")).toEqual([]); + }); + + it("starts from chunk ids in a search result before the answer exists", () => { + const chunkId = "pipelines/productstrategycy27/stages/crawl/17dPGPaiFp-eBYOZVtosqtDa03opeliamHhafZVrRVxU.pdf#28"; + const content = `Found 1 chunks for query in pipeline "productstrategycy27":\n[{"chunk_id":"${chunkId}","chunk_text":"RHEL"}]`; + expect(citationTargetsInChunkResult("get_chunks_for_embeddings", content)).toEqual([ + { pipeline: "productstrategycy27", chunkId }, + ]); + expect(citationTargetsInChunkResult("execute_sql", "select 1")).toEqual([]); + }); +}); + +describe("documentReferences", () => { + it("reads the file id from a document url", () => { + expect(fileIdFromDocumentUrl("https://docs.google.com/document/d/doc-1")).toBe("doc-1"); + expect(fileIdFromDocumentUrl("https://example.com/doc")).toBeNull(); + }); + + it("reads a bare file id and ignores a relative reference link", () => { + const fileId = "1kay0XsAOelA-mJ1kIWvXC2KAAKonMurqxMNIZ0yL95s"; + expect(fileIdFromCitationHref(fileId)).toBe(fileId); + expect(fileIdFromCitationHref(`https://docs.google.com/document/d/${fileId}`)).toBe(fileId); + expect(fileIdFromCitationHref("./calculate_bookings_flag.md")).toBeNull(); + const cited = "1cNsyB6IoH4rpHM6lKeRdGkgO0Vzvy1f4-RvkRJ_6NmA"; + expect(fileIdFromCitationHref(`pipelines/productstrategycy26/stages/crawl/${cited}.pdf#13`)).toBe(cited); + expect(isInAppHref("./calculate_bookings_flag.md")).toBe(true); + expect(isInAppHref("https://docs.redhat.com/en/documentation/rhel")).toBe(false); + }); + + it("builds a reference from retrieved chunks and the answer title", () => { + const messages = [ + { + type: "ai", + tool_calls: [{ + name: "get_chunks_for_embeddings", + content: `Found 1 chunks for query in pipeline "analyzeproductrhel":\n${JSON.stringify([CHUNK])}\n`, + }], + }, + { + type: "ai", + content: "See [CY26 Product Strategy - RHEL](https://docs.google.com/document/d/doc-1).", + tool_calls: [{ + name: "get_citation_image", + args: { chunk_id: "pipelines/a.pdf#4" }, + }], + }, + ]; + const references = documentReferences(messages); + expect(references).toHaveLength(1); + expect(references[0]?.title).toBe("CY26 Product Strategy - RHEL"); + expect(references[0]?.pages).toEqual([3]); + expect(references[0]?.pipeline).toBe("analyzeproductrhel"); + expect(references[0]?.images).toEqual([]); + expect(references[0]?.passages[0]?.text).toContain("Standardize on security."); + }); + + it("uses the answer label when the link target is only the file id", () => { + const fileId = "1kay0XsAOelA-mJ1kIWvXC2KAAKonMurqxMNIZ0yL95s"; + const messages = [ + { + type: "ai", + content: `See [CY26 Product Strategy - RHEL](${fileId}).`, + tool_calls: [{ + name: "get_chunks_for_embeddings", + content: `Found 1 chunks for query in pipeline "analyzeproductrhel":\n${JSON.stringify([{ + ...CHUNK, + file_id: fileId, + }])}\n`, + }], + }, + ]; + expect(documentReferences(messages)[0]?.title).toBe("CY26 Product Strategy - RHEL"); + }); + + it("keeps documents when the chunk list is cut off", () => { + const fileId = "1gLzn3TWNpjj_cReT5kOXFfUjHdcGK5AP"; + const opened = "1RZg2uzP1WbxumFuRkceZcndgQDOsiSOb"; + const full = JSON.stringify([ + { file_id: fileId, chunk_text: "RHEL 10 considerations", chunk_id: "pipelines/a.pdf#1" }, + { file_id: opened, chunk_text: "upgrade notes" }, + ]); + const cut = full.slice(0, full.indexOf("upgrade") + 4); + const content = `Found 100 chunks for query in pipeline "aifrhel":\n${cut}\n\n[Result shortened from 90000 characters.]`; + const messages = [{ + type: "ai", + content: `See [Considerations](${fileId}) and [Upgrading](${opened}).`, + tool_calls: [ + { name: "get_chunks_for_embeddings", content }, + { name: "get_processed_document", args: { pipeline_name: "aifrhel", file_id: opened } }, + ], + }]; + const references = documentReferences(messages); + expect(references.map((item) => item.fileId)).toEqual([fileId, opened]); + expect(references.every((item) => item.pipeline === "aifrhel")).toBe(true); + expect(references[0]?.title).toBe("Considerations"); + }); + + it("keeps each cited chunk as its own page", () => { + const edge = "1cNsyB6IoH4rpHM6lKeRdGkgO0Vzvy1f4-RvkRJ_6NmA"; + const rhel = "17dPGPaiFp-eBYOZVtosqtDa03opeliamHhafZVrRVxU"; + const edgeChunk = `pipelines/productstrategycy26/stages/crawl/${edge}.pdf#13`; + const rhelChunk = `pipelines/productstrategycy27/stages/crawl/${rhel}.pdf#28`; + const messages = [{ + type: "ai", + content: `See [Edge Strategy, p.9](${edgeChunk}) and [RHEL Strategy, p.19](${rhelChunk}).`, + tool_calls: [{ + name: "get_chunks_for_embeddings", + content: `Found 2 chunks for query in pipeline "productstrategycy26":\n${JSON.stringify([ + { file_id: edge, chunk_id: edgeChunk, chunk_text: "Edge page", regions: [{ pageNo: 9 }] }, + { file_id: rhel, chunk_id: rhelChunk, chunk_text: "RHEL page", regions: [{ pageNo: 19 }] }, + ])}\n`, + }], + }]; + const pages = viewerPages(documentReferences(messages)); + expect(pages.map((page) => page.page)).toEqual([9, 19]); + expect(pages.map((page) => page.chunkId)).toEqual([edgeChunk, rhelChunk]); + }); +}); diff --git a/src/frontend/lib/document-references.ts b/src/frontend/lib/document-references.ts new file mode 100644 index 00000000..23f5e12a --- /dev/null +++ b/src/frontend/lib/document-references.ts @@ -0,0 +1,448 @@ +import type { CitationImage } from "../../server/citation-images"; + +export type DocumentPassage = { + page: number | null; + text: string; + chunkId: string | null; +}; + +export type DocumentReference = { + fileId: string; + title: string; + pipeline: string; + pages: number[]; + passages: DocumentPassage[]; + images: CitationImage[]; +}; + +const MAX_DOCUMENTS = 8; +const MAX_PASSAGES = 8; +const CHUNK_FILE = /(?:^|\/)([A-Za-z0-9_-]{25,})\.pdf(?:#(\d+))?$/i; +const CHUNK_PATH = /pipelines\/([A-Za-z0-9_-]+)\/[^\s)"']+?\.pdf#\d+/g; +const CHUNK_ID_FIELD = /"chunk_id"\s*:\s*"([^"]+)"/g; +const MAX_WARM_PAGES = 12; +const PIPELINE_IN_PATH = /(?:^|\/)pipelines\/([A-Za-z0-9_-]+)\//i; +const PAGE_IN_LABEL = /,\s*p\.?\s*(\d+)\s*$/i; +const MAX_PASSAGE_CHARS = 12000; +const MAX_TITLE_CHARS = 90; + +const DOCUMENT_URL = /https:\/\/(?:docs|drive)\.google\.com\/(?:document|spreadsheets|presentation|file)(?:\/u\/\d+)?\/d\/([A-Za-z0-9_-]+)/; +const MARKDOWN_LINK = /\[([^\]]+)\]\(([^)\s]+)\)/g; +const SPLIT_CITATION = /\[([^[\]]+?)\][ \t\r\n]+\((pipelines\/[^\s)]+?\.pdf(?:#\d+)?|https:\/\/(?:docs|drive)\.google\.com\/[^\s)]+)\)/g; + +/** A model often wraps a long citation so the title and path are not one link. */ +export function joinSplitCitationLinks(text: string): string { + return text.replace(SPLIT_CITATION, (_full, label: string, href: string) => { + const title = label.replace(/\s+/g, " ").trim(); + return `[${title}](${href})`; + }); +} +const BARE_FILE_ID = /^[A-Za-z0-9_-]{25,}$/; +const HAS_SCHEME = /^[a-z][a-z0-9+.-]*:/i; + +type ToolCall = { + name?: string; + args?: unknown; + content?: unknown; +}; + +export function fileIdFromDocumentUrl(href: string): string | null { + const match = href.match(new RegExp(DOCUMENT_URL.source)); + return match?.[1] ?? null; +} + +export type CitationTarget = { + fileId: string; + chunkId: string; + pipeline: string; + page: number | null; +}; + +/** A pipeline chunk path such as pipelines/name/stages/crawl/.pdf#13. */ +export function citationTarget(href: string, label = ""): CitationTarget | null { + const trimmed = href.trim(); + const match = CHUNK_FILE.exec(trimmed); + if (!match?.[1]) return null; + const pageMatch = PAGE_IN_LABEL.exec(label.trim()); + const page = pageMatch?.[1] ? Number(pageMatch[1]) : null; + return { + fileId: match[1], + chunkId: trimmed, + pipeline: PIPELINE_IN_PATH.exec(trimmed)?.[1] ?? "", + page: page != null && Number.isFinite(page) ? page : null, + }; +} + +export type WarmTarget = { pipeline: string; chunkId: string }; + +function pushWarmTarget(targets: WarmTarget[], seen: Set, pipeline: string, chunkId: string): void { + if (!pipeline || !chunkId || seen.has(chunkId) || targets.length >= MAX_WARM_PAGES) return; + seen.add(chunkId); + targets.push({ pipeline, chunkId }); +} + +/** Chunk paths already written into the answer, including a partial stream. */ +export function citationTargetsInText(text: string): WarmTarget[] { + const targets: WarmTarget[] = []; + const seen = new Set(); + for (const match of text.matchAll(new RegExp(CHUNK_PATH.source, "g"))) { + pushWarmTarget(targets, seen, match[1] ?? "", match[0] ?? ""); + } + return targets; +} + +/** Cited pages inside a chunk-search result, so images can start before the answer. */ +export function citationTargetsInChunkResult(name: string | undefined, content: unknown): WarmTarget[] { + const text = textOf(content); + const namedChunkSearch = typeof name === "string" && name.endsWith("get_chunks_for_embeddings"); + if (!namedChunkSearch && !text.includes('"chunk_id"')) return []; + const headerPipeline = /pipeline "([^"]+)"/.exec(text)?.[1] ?? ""; + const targets: WarmTarget[] = []; + const seen = new Set(); + for (const match of text.matchAll(new RegExp(CHUNK_ID_FIELD.source, "g"))) { + const chunkId = match[1] ?? ""; + const pipeline = /pipelines\/([A-Za-z0-9_-]+)\//.exec(chunkId)?.[1] ?? headerPipeline; + pushWarmTarget(targets, seen, pipeline, chunkId); + } + return targets; +} + +/** A document URL, a bare file id, or a pipeline chunk path. */ +export function fileIdFromCitationHref(href: string): string | null { + const fromUrl = fileIdFromDocumentUrl(href); + if (fromUrl) return fromUrl; + const fromChunk = citationTarget(href); + if (fromChunk) return fromChunk.fileId; + const trimmed = href.trim(); + return BARE_FILE_ID.test(trimmed) ? trimmed : null; +} + +/** Relative and same-app links resolve under /chat and open another chat. */ +export function isInAppHref(href: string, origin = ""): boolean { + const trimmed = href.trim(); + if (!trimmed || trimmed.startsWith("#")) return true; + if (!HAS_SCHEME.test(trimmed)) return true; + if (!origin) return false; + try { + return new URL(trimmed).origin === origin; + } catch { + return true; + } +} + +function isChunkSearch(name: string | undefined): boolean { + return typeof name === "string" && name.endsWith("get_chunks_for_embeddings"); +} + +function isProcessedDocument(name: string | undefined): boolean { + return typeof name === "string" && name.endsWith("get_processed_document"); +} + +function argumentText(args: unknown, key: string): string { + if (!args || typeof args !== "object" || Array.isArray(args)) return ""; + const value = (args as Record)[key]; + return typeof value === "string" ? value : ""; +} + +function parseJsonArray(text: string): unknown[] { + const start = text.indexOf("["); + if (start < 0) return []; + let depth = 0; + let inString = false; + let escaped = false; + for (let index = start; index < text.length; index += 1) { + const char = text[index]; + if (inString) { + if (escaped) escaped = false; + else if (char === "\\") escaped = true; + else if (char === '"') inString = false; + continue; + } + if (char === '"') inString = true; + else if (char === "[") depth += 1; + else if (char === "]") { + depth -= 1; + if (depth === 0) { + try { + const value: unknown = JSON.parse(text.slice(start, index + 1)); + return Array.isArray(value) ? value : []; + } catch { + return completeObjects(text); + } + } + } + } + return completeObjects(text); +} + +function completeObjects(text: string): unknown[] { + const start = text.indexOf("["); + if (start < 0) return []; + const objects: unknown[] = []; + let depth = 0; + let objectStart = -1; + let inString = false; + let escaped = false; + for (let index = start + 1; index < text.length; index += 1) { + const char = text[index]; + if (inString) { + if (escaped) escaped = false; + else if (char === "\\") escaped = true; + else if (char === '"') inString = false; + continue; + } + if (char === '"') inString = true; + else if (char === "{") { + if (depth === 0) objectStart = index; + depth += 1; + } else if (char === "}") { + depth -= 1; + if (depth === 0 && objectStart >= 0) { + objects.push(parsedValue(text.slice(objectStart, index + 1))); + objectStart = -1; + } + } else if (char === "]" && depth === 0) break; + } + return objects.filter((item) => item !== undefined); +} + +function parsedValue(text: string): unknown { + try { + return JSON.parse(text) as unknown; + } catch { + return undefined; + } +} + +function titleFromText(text: string): string { + const line = text.split("\n").map((part) => part.trim()).find(Boolean) ?? "Document"; + const plain = line + .replace(/^#+\s*/, "") + .replace(/\*\*/g, "") + .replace(/\[([^\]]+)\]\([^)]+\)/g, "$1"); + if (plain.length <= MAX_TITLE_CHARS) return plain; + return `${plain.slice(0, MAX_TITLE_CHARS - 3)}...`; +} + +function pageNumbers(regions: unknown): number[] { + if (!Array.isArray(regions)) return []; + const pages = new Set(); + for (const region of regions) { + if (!region || typeof region !== "object") continue; + const page = (region as { pageNo?: unknown }).pageNo; + if (typeof page === "number" && Number.isFinite(page)) pages.add(page); + } + return [...pages].sort((left, right) => left - right); +} + +function toolCalls(messages: readonly unknown[]): ToolCall[] { + const calls: ToolCall[] = []; + for (const message of messages) { + if (!message || typeof message !== "object") continue; + const list = (message as { tool_calls?: unknown }).tool_calls; + if (!Array.isArray(list)) continue; + for (const call of list) { + if (call && typeof call === "object") calls.push(call as ToolCall); + } + } + return calls; +} + +function textOf(content: unknown): string { + if (typeof content === "string") return content; + if (!Array.isArray(content)) return ""; + return content.map((block) => { + if (typeof block === "string") return block; + if (block && typeof block === "object" && (block as { type?: unknown }).type === "text") { + const text = (block as { text?: unknown }).text; + return typeof text === "string" ? text : ""; + } + return ""; + }).join(""); +} + +function answerText(messages: readonly unknown[]): string { + const parts: string[] = []; + for (const message of messages) { + if (!message || typeof message !== "object") continue; + if (String((message as { type?: unknown }).type).toLowerCase() !== "ai") continue; + parts.push(textOf((message as { content?: unknown }).content)); + } + return parts.join("\n"); +} + +function titlesFromAnswer(text: string): Map { + const titles = new Map(); + for (const match of joinSplitCitationLinks(text).matchAll(MARKDOWN_LINK)) { + const fileId = fileIdFromCitationHref(match[2] ?? ""); + const title = match[1]?.trim(); + if (fileId && title && !titles.has(fileId)) titles.set(fileId, title); + } + return titles; +} + +function rememberProcessedDocuments( + calls: ToolCall[], + byFile: Map, + titles: Map, +): void { + for (const call of calls) { + if (!isProcessedDocument(call.name)) continue; + const fileId = argumentText(call.args, "file_id"); + const pipeline = argumentText(call.args, "pipeline_name"); + if (!fileId) continue; + const existing = byFile.get(fileId); + if (existing) { + if (!existing.pipeline) existing.pipeline = pipeline; + continue; + } + byFile.set(fileId, emptyReference(fileId, titles.get(fileId) ?? "Document", pipeline)); + } + const pipelines = new Set([...byFile.values()].map((item) => item.pipeline).filter(Boolean)); + if (pipelines.size !== 1) return; + const pipeline = [...pipelines][0] ?? ""; + for (const [fileId, title] of titles) { + if (!byFile.has(fileId)) byFile.set(fileId, emptyReference(fileId, title, pipeline)); + } +} + +function emptyReference(fileId: string, title: string, pipeline: string): DocumentReference { + return { fileId, title, pipeline, pages: [], passages: [], images: [] }; +} + +function rememberPage(reference: DocumentReference, page: number | null): void { + if (page == null || reference.pages.includes(page)) return; + reference.pages.push(page); + reference.pages.sort((left, right) => left - right); +} + +function addCitedPassage(reference: DocumentReference, target: CitationTarget): void { + const found = reference.passages.find((passage) => passage.chunkId === target.chunkId); + if (found) { + if (found.page == null && target.page != null) found.page = target.page; + rememberPage(reference, target.page ?? found.page); + return; + } + if (reference.passages.length >= MAX_PASSAGES) return; + reference.passages.push({ page: target.page, text: "", chunkId: target.chunkId }); + rememberPage(reference, target.page); +} + +function rememberAnswerCitations(text: string, byFile: Map): void { + for (const match of joinSplitCitationLinks(text).matchAll(MARKDOWN_LINK)) { + const label = match[1]?.trim() ?? ""; + const target = citationTarget(match[2] ?? "", label); + if (!target) continue; + const existing = byFile.get(target.fileId); + if (!existing) { + const created = emptyReference(target.fileId, label || "Document", target.pipeline); + addCitedPassage(created, target); + byFile.set(target.fileId, created); + continue; + } + if (!existing.pipeline && target.pipeline) existing.pipeline = target.pipeline; + if (existing.title === "Document" && label) existing.title = label; + addCitedPassage(existing, target); + } +} + +export type ViewerPage = { + fileId: string; + title: string; + pipeline: string; + page: number | null; + chunkId: string | null; + text: string; +}; + +/** Cited pages in reference order. One document can contribute several pages. */ +export function viewerPages(references: readonly DocumentReference[]): ViewerPage[] { + const pages: ViewerPage[] = []; + for (const reference of references) { + if (reference.passages.length === 0) { + pages.push({ + fileId: reference.fileId, + title: reference.title, + pipeline: reference.pipeline, + page: reference.pages[0] ?? null, + chunkId: null, + text: "", + }); + continue; + } + for (const passage of reference.passages) { + pages.push({ + fileId: reference.fileId, + title: reference.title, + pipeline: reference.pipeline, + page: passage.page, + chunkId: passage.chunkId, + text: passage.text, + }); + } + } + return pages; +} + +export function viewerStartIndex(pages: readonly ViewerPage[], fileId: string, href = ""): number { + const chunkId = href.trim(); + if (chunkId) { + const exact = pages.findIndex((page) => page.chunkId === chunkId); + if (exact >= 0) return exact; + } + const first = pages.findIndex((page) => page.fileId === fileId); + return first >= 0 ? first : 0; +} + +/** Documents retrieved in this turn, cited ones first. */ +export function documentReferences(messages: readonly unknown[]): DocumentReference[] { + const titles = titlesFromAnswer(answerText(messages)); + const byFile = new Map(); + + for (const call of toolCalls(messages)) { + const content = textOf(call.content); + if (!isChunkSearch(call.name) || !content) continue; + const pipeline = /pipeline "([^"]+)"/.exec(content)?.[1] ?? ""; + for (const item of parseJsonArray(content)) { + if (!item || typeof item !== "object") continue; + const record = item as { file_id?: unknown; chunk_text?: unknown; chunk_id?: unknown; regions?: unknown }; + const fileId = typeof record.file_id === "string" ? record.file_id : ""; + const text = typeof record.chunk_text === "string" ? record.chunk_text.trim() : ""; + if (!fileId || !text) continue; + const chunkId = typeof record.chunk_id === "string" ? record.chunk_id : null; + const pages = pageNumbers(record.regions); + const passage: DocumentPassage = { + page: pages[0] ?? null, + text: text.length > MAX_PASSAGE_CHARS ? `${text.slice(0, MAX_PASSAGE_CHARS - 3)}...` : text, + chunkId, + }; + const existing = byFile.get(fileId); + if (!existing) { + byFile.set(fileId, { + fileId, + title: titles.get(fileId) ?? titleFromText(text), + pipeline, + pages, + passages: [passage], + images: [], + }); + continue; + } + for (const page of pages) { + if (!existing.pages.includes(page)) existing.pages.push(page); + } + existing.pages.sort((left, right) => left - right); + if (existing.passages.length < MAX_PASSAGES) existing.passages.push(passage); + } + } + + rememberProcessedDocuments(toolCalls(messages), byFile, titles); + rememberAnswerCitations(answerText(messages), byFile); + + const cited = [...titles.keys()].filter((fileId) => byFile.has(fileId)); + const rest = [...byFile.keys()].filter((fileId) => !cited.includes(fileId)); + return [...cited, ...rest].slice(0, MAX_DOCUMENTS).flatMap((fileId) => { + const reference = byFile.get(fileId); + return reference ? [reference] : []; + }); +} diff --git a/src/frontend/lib/processed-document.test.ts b/src/frontend/lib/processed-document.test.ts new file mode 100644 index 00000000..bd2d011a --- /dev/null +++ b/src/frontend/lib/processed-document.test.ts @@ -0,0 +1,89 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { DocumentReference } from "./document-references"; +import { clearCitedPageCache, loadDocumentPages, loadDocumentView, loadPageImages, peekPageImage, processedDocumentMarkdown } from "./processed-document"; + +const reference: DocumentReference = { + fileId: "doc-1", + title: "CY26 RHEL", + pipeline: "analyzeproductrhel", + pages: [], + passages: [{ page: null, text: "short passage", chunkId: "path#1" }], + images: [], +}; + +describe("processedDocumentMarkdown", () => { + beforeEach(() => { + clearCitedPageCache(); + }); + it("reads the document body and decodes entities", () => { + const payload = { + content: [{ + type: "text", + text: `Processed document:\n${JSON.stringify({ markdown_content: "R&D {kept}" })}`, + }], + }; + expect(processedDocumentMarkdown(payload)).toBe("R&D {kept}"); + }); + + it("returns nothing when the tool reports an error", () => { + expect(processedDocumentMarkdown({ isError: true, content: [{ type: "text", text: "{}" }] })).toBe(""); + }); +}); + +describe("loadDocumentView", () => { + beforeEach(() => { + clearCitedPageCache(); + }); + + it("loads the processed document and the page image", async () => { + const callTool = vi.fn(async (_mcp: string, name: string) => { + if (name === "get_processed_document") { + return { content: [{ type: "text", text: `Doc:\n${JSON.stringify({ markdown_content: "## Strategy" })}` }] }; + } + return { content: [{ type: "image", data: "abc", mimeType: "image/png" }] }; + }); + const view = await loadDocumentView(reference, callTool); + expect(view.markdown).toBe("## Strategy"); + expect(view.images).toEqual([{ mimeType: "image/png", data: "abc" }]); + expect(callTool).toHaveBeenCalledWith("unstructured-mcp", "get_citation_image", { + pipeline_name: "analyzeproductrhel", + chunk_id: "path#1", + }); + }); + + it("reports the first page before the later pages return", async () => { + let releaseSecond: () => void = () => undefined; + const second = new Promise((resolve) => { + releaseSecond = resolve; + }); + const callTool = vi.fn(async (_mcp: string, _name: string, args: { chunk_id?: string }) => { + if (args.chunk_id === "b") await second; + return { content: [{ type: "image", data: args.chunk_id, mimeType: "image/png" }] }; + }); + const seen: string[][] = []; + const pending = loadDocumentPages("productstrategycy26", ["a", "b"], [], callTool, (images) => { + seen.push(images.map((image) => image.data)); + }); + await vi.waitFor(() => expect(seen[0]).toEqual(["a"])); + releaseSecond(); + await pending; + expect(seen.at(-1)).toEqual(["a", "b"]); + }); + + it("keeps the retrieved passage when the document call fails", async () => { + const callTool = vi.fn(async () => { + throw new Error("down"); + }); + const view = await loadDocumentView({ ...reference, passages: [{ ...reference.passages[0], chunkId: null }] }, callTool); + expect(view.markdown).toBe("short passage"); + expect(view.images).toEqual([]); + }); + + it("reuses a page image that was already fetched", async () => { + const callTool = vi.fn(async () => ({ content: [{ type: "image", data: "cached", mimeType: "image/png" }] })); + await loadPageImages("productstrategycy26", "chunk-9", [], callTool); + expect(peekPageImage("productstrategycy26", "chunk-9")?.data).toBe("cached"); + await loadPageImages("productstrategycy26", "chunk-9", [], callTool); + expect(callTool).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/frontend/lib/processed-document.ts b/src/frontend/lib/processed-document.ts new file mode 100644 index 00000000..838eb4ee --- /dev/null +++ b/src/frontend/lib/processed-document.ts @@ -0,0 +1,216 @@ +import { citationImagesFromContent, type CitationImage } from "../../server/citation-images"; +import type { DocumentReference } from "./document-references"; + +const MCP_SERVER = "unstructured-mcp"; +const MAX_DOCUMENT_CHARS = 500_000; + +export type McpToolCaller = ( + mcpName: string, + name: string, + args: Record, +) => Promise<{ content?: unknown; isError?: boolean }>; + +export type DocumentView = { + markdown: string; + images: CitationImage[]; +}; + +function textOfPayload(payload: unknown): string { + if (typeof payload === "string") return payload; + if (!payload || typeof payload !== "object") return ""; + const content = (payload as { content?: unknown }).content; + if (typeof content === "string") return content; + if (!Array.isArray(content)) return ""; + return content.map((block) => { + if (!block || typeof block !== "object") return ""; + if ((block as { type?: unknown }).type !== "text") return ""; + const text = (block as { text?: unknown }).text; + return typeof text === "string" ? text : ""; + }).join("\n"); +} + +function objectFromText(text: string): Record | null { + const start = text.indexOf("{"); + if (start < 0) return null; + let depth = 0; + let inString = false; + let escaped = false; + for (let index = start; index < text.length; index += 1) { + const char = text[index]; + if (inString) { + if (escaped) escaped = false; + else if (char === "\\") escaped = true; + else if (char === '"') inString = false; + continue; + } + if (char === '"') inString = true; + else if (char === "{") depth += 1; + else if (char === "}") { + depth -= 1; + if (depth === 0) return parsedObject(text.slice(start, index + 1)); + } + } + return null; +} + +function parsedObject(text: string): Record | null { + try { + const value: unknown = JSON.parse(text); + if (!value || typeof value !== "object" || Array.isArray(value)) return null; + return value as Record; + } catch { + return null; + } +} + +function decodeEntities(value: string): string { + return value + .replace(/&/g, "&") + .replace(/</g, "<") + .replace(/>/g, ">") + .replace(/"/g, '"') + .replace(/'/g, "'"); +} + +/** Markdown body of a get_processed_document tool result. */ +export function processedDocumentMarkdown(payload: unknown): string { + if (payload && typeof payload === "object" && (payload as { isError?: unknown }).isError === true) return ""; + const record = objectFromText(textOfPayload(payload)); + const markdown = record?.markdown_content; + if (typeof markdown !== "string") return ""; + return decodeEntities(markdown).slice(0, MAX_DOCUMENT_CHARS); +} + +const inflightMarkdown = new Map>(); +const inflightImages = new Map>(); +const storedImages = new Map(); + +function imageKey(pipeline: string, chunkId: string): string { + return `${pipeline}:${chunkId}`; +} + +/** The page image, if this session already fetched it. `undefined` means it has not finished. */ +export function peekPageImage(pipeline: string, chunkId: string): CitationImage | null | undefined { + const stored = storedImages.get(imageKey(pipeline, chunkId)); + if (!stored) return undefined; + return stored[0] ?? null; +} + +export function clearCitedPageCache(): void { + inflightImages.clear(); + storedImages.clear(); +} + +async function readProcessed(pipeline: string, fileId: string, callTool: McpToolCaller): Promise { + try { + const result = await callTool(MCP_SERVER, "get_processed_document", { + pipeline_name: pipeline, + file_id: fileId, + }); + return processedDocumentMarkdown(result); + } catch { + return ""; + } +} + +async function readImages(pipeline: string, chunkId: string, callTool: McpToolCaller): Promise { + try { + const result = await callTool(MCP_SERVER, "get_citation_image", { + pipeline_name: pipeline, + chunk_id: chunkId, + }); + return citationImagesFromContent(result.content); + } catch { + return []; + } +} + +/** One in-flight fetch per document, shared if the viewer mounts twice. */ +export function loadProcessedMarkdown(pipeline: string, fileId: string, callTool: McpToolCaller): Promise { + if (!pipeline || !fileId) return Promise.resolve(""); + const key = `${pipeline}:${fileId}`; + const existing = inflightMarkdown.get(key); + if (existing) return existing; + const pending = readProcessed(pipeline, fileId, callTool).finally(() => { + inflightMarkdown.delete(key); + }); + inflightMarkdown.set(key, pending); + return pending; +} + +/** One fetch per chunk for the whole session. A later open reuses the stored image. */ +export function loadPageImages( + pipeline: string, + chunkId: string, + existing: CitationImage[], + callTool: McpToolCaller, +): Promise { + if (existing.length > 0) return Promise.resolve(existing); + if (!pipeline || !chunkId) return Promise.resolve([]); + const key = imageKey(pipeline, chunkId); + const stored = storedImages.get(key); + if (stored) return Promise.resolve(stored); + const current = inflightImages.get(key); + if (current) return current; + const pending = readImages(pipeline, chunkId, callTool).then((images) => { + storedImages.set(key, images); + return images; + }).finally(() => { + inflightImages.delete(key); + }); + inflightImages.set(key, pending); + return pending; +} + +/** Start every cited page image while the answer is on screen, before a click. */ +export function prefetchCitedPages( + pages: readonly { pipeline: string; chunkId: string | null }[], + callTool: McpToolCaller, +): void { + for (const page of pages) { + if (!page.pipeline || !page.chunkId) continue; + void loadPageImages(page.pipeline, page.chunkId, [], callTool); + } +} + +function takePages(pages: CitationImage[][], limit = 6): CitationImage[] { + const images: CitationImage[] = []; + for (const page of pages) { + for (const image of page) { + images.push(image); + if (images.length >= limit) return images; + } + } + return images; +} + +/** Page pictures in citation order. The first page is reported before the later pages return. */ +export async function loadDocumentPages( + pipeline: string, + chunkIds: readonly string[], + existing: CitationImage[], + callTool: McpToolCaller, + onUpdate?: (images: CitationImage[]) => void, +): Promise { + if (existing.length > 0) return existing; + const ids = [...new Set(chunkIds.filter((id) => id.length > 0))].slice(0, 3); + if (ids.length === 0) return []; + const first = await loadPageImages(pipeline, ids[0], [], callTool); + let images = takePages([first]); + if (images.length > 0) onUpdate?.(images); + const rest = await Promise.all(ids.slice(1).map((id) => loadPageImages(pipeline, id, [], callTool))); + images = takePages([first, ...rest]); + if (rest.length > 0 && images.length > 0) onUpdate?.(images); + return images; +} + +/** Full document text plus any page image the unstructured server can return. */ +export async function loadDocumentView(reference: DocumentReference, callTool: McpToolCaller): Promise { + const chunkId = reference.passages.find((passage) => passage.chunkId)?.chunkId ?? ""; + const [markdown, images] = await Promise.all([ + loadProcessedMarkdown(reference.pipeline, reference.fileId, callTool), + loadPageImages(reference.pipeline, chunkId, reference.images, callTool), + ]); + const fallback = reference.passages.map((passage) => passage.text).join("\n\n"); + return { markdown: markdown || fallback, images }; +} diff --git a/src/server/citation-images.test.ts b/src/server/citation-images.test.ts new file mode 100644 index 00000000..a0f97d7d --- /dev/null +++ b/src/server/citation-images.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, it } from "vitest"; +import { citationImagesFromContent } from "./citation-images"; + +describe("citationImagesFromContent", () => { + it("keeps png and jpeg blocks and drops text", () => { + const images = citationImagesFromContent([ + { type: "text", text: "Citation for chunk, 1 page(s): 94" }, + { type: "image", base64: "abc", mime_type: "image/png" }, + { type: "image", data: "def", mimeType: "image/jpeg" }, + ]); + expect(images).toEqual([ + { mimeType: "image/png", data: "abc" }, + { mimeType: "image/jpeg", data: "def" }, + ]); + }); + + it("returns nothing for a plain string", () => { + expect(citationImagesFromContent("Citation for chunk")).toEqual([]); + }); + + it("rejects a non-image mime and a remote url", () => { + const images = citationImagesFromContent([ + { type: "image", data: "abc", mimeType: "image/svg+xml" }, + { type: "image", data: "https://example.com/a.png", mimeType: "image/png" }, + ]); + expect(images).toEqual([]); + }); +}); diff --git a/src/server/citation-images.ts b/src/server/citation-images.ts new file mode 100644 index 00000000..5fb93725 --- /dev/null +++ b/src/server/citation-images.ts @@ -0,0 +1,33 @@ +/** A page crop returned by get_citation_image. */ +export type CitationImage = { + mimeType: string; + data: string; +}; + +const MAX_IMAGES = 4; +const MAX_DATA_CHARS = 1_500_000; +const ALLOWED_MIME = new Set(["image/png", "image/jpeg", "image/webp", "image/gif"]); + +function imageBlock(block: unknown): CitationImage | null { + if (!block || typeof block !== "object") return null; + const record = block as Record; + if (record.type !== "image") return null; + const data = typeof record.base64 === "string" ? record.base64 : typeof record.data === "string" ? record.data : ""; + if (!data || data.startsWith("http") || data.length > MAX_DATA_CHARS) return null; + const mime = typeof record.mime_type === "string" ? record.mime_type : typeof record.mimeType === "string" ? record.mimeType : "image/png"; + if (!ALLOWED_MIME.has(mime)) return null; + return { mimeType: mime, data }; +} + +/** Pull page crops out of a tool message. Text blocks are left for the caller. */ +export function citationImagesFromContent(content: unknown): CitationImage[] { + if (!Array.isArray(content)) return []; + const images: CitationImage[] = []; + for (const block of content) { + const image = imageBlock(block); + if (!image) continue; + images.push(image); + if (images.length >= MAX_IMAGES) break; + } + return images; +} From fdd7ec0c11020aba9f092784717d6df7aa29834c Mon Sep 17 00:00:00 2001 From: mimran-khan Date: Thu, 12 Nov 2026 12:32:32 +0530 Subject: [PATCH 2/2] fix: finish interrupt streams so auto-approve can complete An interrupt with no answer text was treated as a dropped stream, so the run stayed loading and the approval never resumed. A paused interrupt now ends the stream, and the auto-approve test reports the thread as interrupted. Signed-off-by: mimran-khan --- e2e/hitl/auto-approve.spec.ts | 17 ++++++++++++++--- src/frontend/hooks/useStreamingAPI.ts | 14 ++++++++++++++ src/frontend/redux/slices/chats.test.ts | 2 +- 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/e2e/hitl/auto-approve.spec.ts b/e2e/hitl/auto-approve.spec.ts index 3670080d..37ab4657 100644 --- a/e2e/hitl/auto-approve.spec.ts +++ b/e2e/hitl/auto-approve.spec.ts @@ -35,6 +35,18 @@ async function makeInterruptStream( callCount: { n: number }, page: import('@playwright/test').Page, ): Promise { + await page.route('**/api/proxy/agent/threads/*', (route) => { + const path = new URL(route.request().url()).pathname; + const isThreadRecord = /\/threads\/(?!search$)[^/]+$/.test(path); + if (route.request().method() === 'GET' && isThreadRecord) { + return route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ status: 'interrupted' }), + }); + } + return route.fallback(); + }); await page.route('**/api/proxy/agent/v1/stream', (route) => { callCount.n++; if (callCount.n === 1) { @@ -78,8 +90,7 @@ test.describe('HITL — auto-approve', () => { const chat = new ChatPage(page); await chat.expectChatRoute(); - // The response should complete without the interrupt banner ever appearing - await chat.waitForAIResponse(15_000); + await expect(page.getByText('Done automatically.')).toBeVisible({ timeout: 15_000 }); await expect(page.getByText('Action Required')).not.toBeVisible(); await expect(page.getByText('Input Required')).not.toBeVisible(); expect(callCount.n).toBeGreaterThanOrEqual(2); @@ -129,7 +140,7 @@ test.describe('HITL — auto-approve', () => { const chat = new ChatPage(page); await chat.expectChatRoute(); - await chat.waitForAIResponse(15_000); + await expect(page.getByText('Done automatically.')).toBeVisible({ timeout: 15_000 }); await expect(page.getByText('Action Required')).not.toBeVisible(); expect(callCount.n).toBeGreaterThanOrEqual(2); }, diff --git a/src/frontend/hooks/useStreamingAPI.ts b/src/frontend/hooks/useStreamingAPI.ts index e2734cc4..aaff6b12 100644 --- a/src/frontend/hooks/useStreamingAPI.ts +++ b/src/frontend/hooks/useStreamingAPI.ts @@ -1004,6 +1004,20 @@ export function useStreamingAPI(threadId: string) { dispatch(resolveAllPendingToolCalls({ chatId: threadId })); } lastSuccessfulConnectionRef.current = Date.now(); + if (streamEndedWithInterruptRef.current) { + dispatch(updateStreamingState({ + chatId: threadId, + state: { + isLoading: false, + isConnected: false, + isReconnecting: false, + reconnectAttempt: 0, + streamDroppedMidResponse: false, + }, + })); + finish('success'); + return; + } const gotFullResponse = messagesRef.current.length > initialMessageCount + 1; if (gotFullResponse && recoveryIntervalRef.current) { clearInterval(recoveryIntervalRef.current); diff --git a/src/frontend/redux/slices/chats.test.ts b/src/frontend/redux/slices/chats.test.ts index b0c99538..4cecbc04 100644 --- a/src/frontend/redux/slices/chats.test.ts +++ b/src/frontend/redux/slices/chats.test.ts @@ -168,7 +168,7 @@ describe('chats slice — appendMessageToChat', () => { describe('chats slice — subagent transcript', () => { function withSubagentCard() { - let s = chatsReducer(initialState(), addChat(makeChat('c1'))); + const s = chatsReducer(initialState(), addChat(makeChat('c1'))); const card: TestMessage = { type: 'ai', content: 'Asking the analyst.',