diff --git a/apps/web/components/chat/message/agent-message.tsx b/apps/web/components/chat/message/agent-message.tsx
index 7e0df6bc..64a59d5f 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 7abe6859..5bd6f0fa 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