fix(chat): align source fallback gate with rendered runs

This commit is contained in:
Dhravya 2026-06-27 04:54:12 +00:00
parent 6798b565cb
commit 865962ae19
3 changed files with 77 additions and 22 deletions

View file

@ -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 (
<div
key={`${message.id}-${partIndex}`}

View file

@ -1,7 +1,9 @@
import { describe, expect, it } from "bun:test"
import {
hasRenderableSourceAnnotations,
isSafeSourceId,
parseSourceAnnotatedMarkdown,
sourceAnnotatedTextRun,
stripSourceMarkup,
} from "./source-annotations"
@ -78,6 +80,31 @@ describe("source annotation parsing", () => {
)
})
it("checks inline annotations against rendered text runs", () => {
const allowedSourceIds = new Set(["S1"])
expect(
hasRenderableSourceAnnotations(
[
{ type: "text", text: 'Lead <response source="S1">supported' },
{ type: "tool-recallContext" },
{ type: "text", text: " claim</response>" },
],
allowedSourceIds,
),
).toBe(false)
const parts = [
{ type: "text", text: 'Lead <response source="S1">supported' },
{ type: "source-url", sourceId: "web", url: "https://example.com" },
{ type: "text", text: " claim</response>" },
]
expect(sourceAnnotatedTextRun(parts, 0)).toBe(
'Lead <response source="S1">supported claim</response>',
)
expect(sourceAnnotatedTextRun(parts, 2)).toBeNull()
expect(hasRenderableSourceAnnotations(parts, allowedSourceIds)).toBe(true)
})
it("strips source markup for copy text", () => {
expect(
stripSourceMarkup('Alpha <response source="S1">Beta</response>'),

View file

@ -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<string>,
): 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