From 6798b565cb5d698a59f74fc5221a5d9d1fcb9559 Mon Sep 17 00:00:00 2001 From: Dhravya <63950637+Dhravya@users.noreply.github.com> Date: Sat, 27 Jun 2026 04:35:28 +0000 Subject: [PATCH 1/2] fix(chat): show memory source fallbacks --- .../components/chat/message/agent-message.tsx | 111 ++++++++++++------ apps/web/lib/chat-memory-tools.test.ts | 93 +++++++++++++++ apps/web/lib/chat-memory-tools.ts | 64 +++++++++- 3 files changed, 227 insertions(+), 41 deletions(-) diff --git a/apps/web/components/chat/message/agent-message.tsx b/apps/web/components/chat/message/agent-message.tsx index eac0f7b9b..7e0df6bc8 100644 --- a/apps/web/components/chat/message/agent-message.tsx +++ b/apps/web/components/chat/message/agent-message.tsx @@ -23,6 +23,7 @@ import { isWebSearchToolName } from "@/lib/chat-web-search-tools" import { buildCitationIndex, fetchDocumentsByIds, + getCitationDisplay, getDocumentSourceUrl, isMemoryToolOutputReady, mapDocumentsByKnownIds, @@ -190,37 +191,64 @@ function CitationLink({ ) } -function sourceTitle( +function documentForCitationTarget( target: CitationTarget, - document?: DocumentWithMemories, -): string { + documentByKnownId: Map, +): DocumentWithMemories | undefined { return ( - document?.title?.trim() || - target.title?.trim() || - document?.customId || - target.customId || - target.documentId || - target.sourceId + (target.documentId + ? documentByKnownId.get(target.documentId) + : undefined) ?? + (target.customId ? documentByKnownId.get(target.customId) : undefined) ) } -function sourceSummary( - target: CitationTarget, - document?: DocumentWithMemories, -): string | null { - const summary = - document?.summary || - target.summary || - (document as { content?: string } | undefined)?.content || - null - return summary ? summary.trim() : null -} +function MemorySourcesFallback({ + citationIndex, + documentByKnownId, +}: { + citationIndex: Map + documentByKnownId: Map +}) { + const sources = Array.from(citationIndex.values()).slice(0, 6) + if (sources.length === 0) return null -function sourceKind( - target: CitationTarget, - document?: DocumentWithMemories, -): string { - return (document?.type || target.type || "memory").replaceAll("_", " ") + return ( +
+ Sources + {sources.map((target) => { + const document = documentForCitationTarget(target, documentByKnownId) + const display = getCitationDisplay(target, document) + const url = safeExternalUrl( + document ? getDocumentSourceUrl(document) : target.url, + ) + const label = `${target.sourceId}: ${display.title}` + const className = + "max-w-56 truncate rounded-full border border-white/10 bg-white/[0.035] px-2 py-1 text-white/55 transition-colors hover:bg-white/[0.06] hover:text-white/75" + + return url ? ( + + {label} + + ) : ( + + {label} + + ) + })} +
+ ) } function SourceCitationLink({ @@ -237,16 +265,11 @@ function SourceCitationLink({ const target = citationIndex.get(sourceId) if (!target) return <>{children} - const document = - (target.documentId - ? documentByKnownId.get(target.documentId) - : undefined) ?? - (target.customId ? documentByKnownId.get(target.customId) : undefined) + const document = documentForCitationTarget(target, documentByKnownId) const url = safeExternalUrl( document ? getDocumentSourceUrl(document) : target.url, ) - const title = sourceTitle(target, document) - const summary = sourceSummary(target, document) + const display = getCitationDisplay(target, document) return ( @@ -274,15 +297,15 @@ function SourceCitationLink({ - {title} + {display.title} - {sourceKind(target, document)} + {display.kind} - {summary ? ( + {display.summary ? ( - {summary} + {display.summary} ) : null} {url ? ( @@ -754,6 +777,16 @@ export function AgentMessage({ () => makeMarkdownComponents(webSources, citationIndex, documentByKnownId), [webSources, citationIndex, documentByKnownId], ) + const hasInlineSourceAnnotations = useMemo( + () => + parseSourceAnnotatedMarkdown( + messageText, + allowedSourceIds, + ).markdown.includes("#sm-source:"), + [messageText, allowedSourceIds], + ) + const showMemorySourcesFallback = + citationIndex.size > 0 && !hasInlineSourceAnnotations const responseModelLabel = responseModel ? `${modelNames[responseModel].name} ${modelNames[responseModel].version}` : null @@ -865,6 +898,12 @@ export function AgentMessage({ } return null })} + {showMemorySourcesFallback && ( + + )} {hasAssistantText && ( diff --git a/apps/web/lib/chat-memory-tools.test.ts b/apps/web/lib/chat-memory-tools.test.ts index 9b6411c7a..f7e8a4c95 100644 --- a/apps/web/lib/chat-memory-tools.test.ts +++ b/apps/web/lib/chat-memory-tools.test.ts @@ -4,6 +4,7 @@ import { buildCitationIndex, extractDocumentIdsFromMemoryOutput, extractMemoryToolOutputs, + getCitationDisplay, getDocumentSourceUrl, mapDocumentsByKnownIds, } from "./chat-memory-tools" @@ -76,9 +77,101 @@ describe("chat memory tool citation mapping", () => { expect(index.get("S1")?.documentId).toBe("docA") expect(index.get("S1")?.customId).toBe("customA") + expect(index.get("S1")?.content).toBe("memo") expect(index.has("ignored")).toBe(false) }) + it("maps source ids to matching result ids when citation ids are absent", () => { + const outputs = extractMemoryToolOutputs({ + parts: [ + { + type: "tool-searchMemories", + state: "output-available", + output: { + sourceIds: ["memory_no_citation"], + results: [ + { + id: "memory_no_citation", + kind: "memory", + content: "Fallback source text.", + }, + ], + }, + }, + ], + }) + + expect(buildCitationIndex(outputs).get("memory_no_citation")).toMatchObject( + { + sourceId: "memory_no_citation", + memoryId: "memory_no_citation", + content: "Fallback source text.", + }, + ) + }) + + it("keeps memory text for citations without source documents", () => { + const outputs = extractMemoryToolOutputs({ + parts: [ + { + type: "tool-recallContext", + state: "output-available", + output: { + sourceIds: ["memory_1"], + results: [ + { + id: "memory_1", + citationId: "memory_1", + kind: "memory", + content: "User prefers concise answers.", + }, + ], + }, + }, + ], + }) + + expect(buildCitationIndex(outputs).get("memory_1")).toMatchObject({ + sourceId: "memory_1", + memoryId: "memory_1", + kind: "memory", + content: "User prefers concise answers.", + }) + }) + + it("uses memory display only for citations without document metadata", () => { + expect( + getCitationDisplay({ + sourceId: "memory_1", + memoryId: "memory_1", + kind: "memory", + content: "User prefers concise answers.", + }), + ).toEqual({ + title: "Memory", + kind: "memory", + summary: "User prefers concise answers.", + }) + }) + + it("prefers tool-provided document metadata over generic memory display", () => { + expect( + getCitationDisplay({ + sourceId: "S1", + memoryId: "memory_1", + content: "Memory text", + documentId: "doc_1", + customId: "custom_1", + title: "Project Plan", + type: "google_doc", + }), + ).toMatchObject({ + title: "Project Plan", + kind: "google doc", + summary: "Memory text", + }) + }) + it("extracts graph highlight document ids from memory outputs", () => { const [output] = extractMemoryToolOutputs(assistantMessage) diff --git a/apps/web/lib/chat-memory-tools.ts b/apps/web/lib/chat-memory-tools.ts index 140c2ae6f..6a39d2f5b 100644 --- a/apps/web/lib/chat-memory-tools.ts +++ b/apps/web/lib/chat-memory-tools.ts @@ -55,6 +55,9 @@ export type MemoryToolOutput = { } export type CitationTarget = { sourceId: string + memoryId?: string | undefined + kind?: string | undefined + content?: string | undefined documentId?: string | undefined customId?: string | null | undefined title?: string | null | undefined @@ -63,6 +66,12 @@ export type CitationTarget = { url?: string | null | undefined } +export type CitationDisplay = { + title: string + summary: string | null + kind: string +} + export type DocumentWithMemories = z.infer< typeof DocumentsWithMemoriesResponseSchema >["documents"][0] @@ -156,10 +165,15 @@ function citationTargetForResult( doc?.id ?? result.documentIds?.find(Boolean) ?? result.documentId - const target = documentTarget(sourceId, doc) - target.documentId = target.documentId ?? docId - target.customId = target.customId ?? result.customId - return target + const docTarget = documentTarget(sourceId, doc) + return { + ...docTarget, + memoryId: result.id, + kind: result.kind, + content: result.content, + documentId: docTarget.documentId ?? docId, + customId: docTarget.customId ?? result.customId, + } } function documentTarget( @@ -247,7 +261,7 @@ export function buildCitationIndex( for (const sourceId of output.sourceIds ?? []) { if (index.has(sourceId)) continue const matchingResult = (output.results ?? []).find( - (result) => result.citationId === sourceId, + (result) => result.citationId === sourceId || result.id === sourceId, ) if (matchingResult) addTarget( @@ -344,6 +358,46 @@ export function mapDocumentsByKnownIds( return map } +function hasDisplayMetadata(target: CitationTarget): boolean { + return !!( + target.title?.trim() || + target.documentId || + target.customId || + target.url || + target.type || + target.summary?.trim() + ) +} + +export function getCitationDisplay( + target: CitationTarget, + document?: DocumentWithMemories, +): CitationDisplay { + const memoryOnly = !document && !hasDisplayMetadata(target) + const summary = + document?.summary || + target.summary || + target.content || + (document as { content?: string } | undefined)?.content || + null + + return { + title: memoryOnly + ? "Memory" + : document?.title?.trim() || + target.title?.trim() || + document?.customId || + target.customId || + target.documentId || + target.sourceId, + summary: summary ? summary.trim() : null, + kind: (document?.type || target.type || target.kind || "memory").replaceAll( + "_", + " ", + ), + } +} + export function getDocumentSourceUrl( document: Pick & { customId?: string | null From 865962ae195548f2ec06babe3ce0e81bd2dfede5 Mon Sep 17 00:00:00 2001 From: Dhravya <63950637+Dhravya@users.noreply.github.com> Date: Sat, 27 Jun 2026 04:54:12 +0000 Subject: [PATCH 2/2] fix(chat): align source fallback gate with rendered runs --- .../components/chat/message/agent-message.tsx | 30 ++++--------- apps/web/lib/source-annotations.test.ts | 27 ++++++++++++ apps/web/lib/source-annotations.ts | 42 +++++++++++++++++++ 3 files changed, 77 insertions(+), 22 deletions(-) diff --git a/apps/web/components/chat/message/agent-message.tsx b/apps/web/components/chat/message/agent-message.tsx index 7e0df6bc8..64a59d5fd 100644 --- a/apps/web/components/chat/message/agent-message.tsx +++ b/apps/web/components/chat/message/agent-message.tsx @@ -32,7 +32,9 @@ import { extractMemoryToolOutputs, } from "@/lib/chat-memory-tools" import { + hasRenderableSourceAnnotations, parseSourceAnnotatedMarkdown, + sourceAnnotatedTextRun, stripSourceMarkup, } from "@/lib/source-annotations" import { modelNames, type ModelId } from "@/lib/models" @@ -778,12 +780,8 @@ export function AgentMessage({ [webSources, citationIndex, documentByKnownId], ) const hasInlineSourceAnnotations = useMemo( - () => - parseSourceAnnotatedMarkdown( - messageText, - allowedSourceIds, - ).markdown.includes("#sm-source:"), - [messageText, allowedSourceIds], + () => hasRenderableSourceAnnotations(message.parts, allowedSourceIds), + [message.parts, allowedSourceIds], ) const showMemorySourcesFallback = citationIndex.size > 0 && !hasInlineSourceAnnotations @@ -831,22 +829,10 @@ export function AgentMessage({ ) } if (part.type === "text") { - // Skip fragments mid-run — source-url citations split one answer into - // many text parts; rendering each separately tears markdown (lists etc.). - let prev = partIndex - 1 - while (prev >= 0 && message.parts[prev]?.type === "source-url") { - prev-- - } - if (prev >= 0 && message.parts[prev]?.type === "text") { - return null - } - let runText = "" - for (let j = partIndex; j < message.parts.length; j++) { - const p = message.parts[j] - if (p?.type === "text") runText += p.text - else if (p?.type === "source-url") continue - else break - } + // source-url citations split one answer into many text parts; + // render each contiguous text/source-url run as one markdown block. + const runText = sourceAnnotatedTextRun(message.parts, partIndex) + if (runText === null) return null return (
{ ) }) + it("checks inline annotations against rendered text runs", () => { + const allowedSourceIds = new Set(["S1"]) + expect( + hasRenderableSourceAnnotations( + [ + { type: "text", text: 'Lead supported' }, + { type: "tool-recallContext" }, + { type: "text", text: " claim" }, + ], + allowedSourceIds, + ), + ).toBe(false) + + const parts = [ + { type: "text", text: 'Lead supported' }, + { type: "source-url", sourceId: "web", url: "https://example.com" }, + { type: "text", text: " claim" }, + ] + expect(sourceAnnotatedTextRun(parts, 0)).toBe( + 'Lead supported claim', + ) + expect(sourceAnnotatedTextRun(parts, 2)).toBeNull() + expect(hasRenderableSourceAnnotations(parts, allowedSourceIds)).toBe(true) + }) + it("strips source markup for copy text", () => { expect( stripSourceMarkup('Alpha Beta'), diff --git a/apps/web/lib/source-annotations.ts b/apps/web/lib/source-annotations.ts index 7abe68595..5bd6f0fa3 100644 --- a/apps/web/lib/source-annotations.ts +++ b/apps/web/lib/source-annotations.ts @@ -169,6 +169,48 @@ export function parseSourceAnnotatedMarkdown( return { markdown: output.join("") } } +export type SourceAnnotationMessagePart = { + type: string + text?: string | undefined +} + +export function sourceAnnotatedTextRun( + parts: readonly SourceAnnotationMessagePart[], + partIndex: number, +): string | null { + const part = parts[partIndex] + if (part?.type !== "text") return null + + let prev = partIndex - 1 + while (prev >= 0 && parts[prev]?.type === "source-url") prev-- + if (prev >= 0 && parts[prev]?.type === "text") return null + + let runText = "" + for (let index = partIndex; index < parts.length; index++) { + const current = parts[index] + if (current?.type === "text") runText += current.text ?? "" + else if (current?.type === "source-url") continue + else break + } + + return runText +} + +export function hasRenderableSourceAnnotations( + parts: readonly SourceAnnotationMessagePart[], + allowedSourceIds: ReadonlySet, +): boolean { + return parts.some((part, index) => { + if (part.type !== "text") return false + const runText = sourceAnnotatedTextRun(parts, index) + if (!runText) return false + return parseSourceAnnotatedMarkdown( + runText, + allowedSourceIds, + ).markdown.includes("#sm-source:") + }) +} + export function stripSourceMarkup(text: string): string { let output = "" let i = 0