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] 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 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 (
 								<div
 									key={`${message.id}-${partIndex}`}
diff --git a/apps/web/lib/source-annotations.test.ts b/apps/web/lib/source-annotations.test.ts
index dce37e1a..65986290 100644
--- a/apps/web/lib/source-annotations.test.ts
+++ b/apps/web/lib/source-annotations.test.ts
@@ -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>'),
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<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