Fix chat source review issues

apps/web/components/chat/message/agent-message.tsx

@@ -24,6 +24,7 @@ import {
 	buildCitationIndex,
 	fetchDocumentsByIds,
 	getDocumentSourceUrl,
+	isMemoryToolOutputReady,
 	mapDocumentsByKnownIds,
 	type CitationTarget,
 	type DocumentWithMemories,
@@ -57,7 +58,7 @@ const TOOL_META: Record<string, { label: string; icon: typeof SearchIcon }> = {
 
 type ToolCallDisplayPart = {
 	type: string
-	state: string
+	state?: string
 	input?: unknown
 	output?: unknown
 	toolCallId?: string
@@ -581,7 +582,9 @@ function ToolCallDisplay({ part }: { part: ToolCallDisplayPart }) {
 	const isDone = part.state === "output-available"
 	const isError = part.state === "error" || part.state === "output-error"
 	const errorText = part.errorText
-	if (isMemoryRetrievalToolName(toolName) && isDone) return null
+	if (isMemoryRetrievalToolName(toolName) && isMemoryToolOutputReady(part)) {
+		return null
+	}
 
 	return (
 		<div className="rounded-lg border border-[#1E2128] bg-[#0D121A] text-xs my-1 overflow-hidden">
@@ -830,7 +833,7 @@ export function AgentMessage({
 								type: "dynamic-tool"
 								toolName: string
 								toolCallId: string
-								state: string
+								state?: string
 								input?: unknown
 								output?: unknown
 								errorText?: string

apps/web/lib/chat-memory-tools.test.ts

@@ -4,6 +4,7 @@ import {
 	buildCitationIndex,
 	extractDocumentIdsFromMemoryOutput,
 	extractMemoryToolOutputs,
+	getDocumentSourceUrl,
 	mapDocumentsByKnownIds,
 } from "./chat-memory-tools"
 
@@ -84,9 +85,9 @@ describe("chat memory tool citation mapping", () => {
 		expect(output && extractDocumentIdsFromMemoryOutput(output.output)).toEqual(
 			["topDoc", "docA", "customA"],
 		)
-		expect(extractHighlightDocumentIdsFromMessages([assistantMessage])).toEqual(
-			["topDoc", "docA", "customA"],
-		)
+		expect(
+			extractHighlightDocumentIdsFromMessages([assistantMessage as never]),
+		).toEqual(["topDoc", "docA", "customA"])
 	})
 
 	it("keeps graph highlights for legacy memory tool states and ids", () => {
@@ -107,10 +108,9 @@ describe("chat memory tool citation mapping", () => {
 		} as const
 
 		expect(extractMemoryToolOutputs(legacyMessage)).toHaveLength(2)
-		expect(extractHighlightDocumentIdsFromMessages([legacyMessage])).toEqual([
-			"legacyDoc",
-			"statelessDoc",
-		])
+		expect(
+			extractHighlightDocumentIdsFromMessages([legacyMessage as never]),
+		).toEqual(["legacyDoc", "statelessDoc"])
 	})
 
 	it("normalizes nested discoverSpaces memory results", () => {
@@ -139,6 +139,34 @@ describe("chat memory tool citation mapping", () => {
 		).toEqual(["spaceDoc", "nestedDoc"])
 	})
 
+	it("builds editable Google source URLs from custom ids and API URLs", () => {
+		expect(
+			getDocumentSourceUrl({
+				type: "google_doc",
+				customId: "docCustom",
+				url: "https://docs.googleapis.com/v1/documents/apiDoc",
+			} as never),
+		).toBe("https://docs.google.com/document/d/docCustom/edit")
+		expect(
+			getDocumentSourceUrl({
+				type: "google_doc",
+				url: "https://docs.googleapis.com/v1/documents/apiDoc",
+			} as never),
+		).toBe("https://docs.google.com/document/d/apiDoc/edit")
+		expect(
+			getDocumentSourceUrl({
+				type: "google_sheet",
+				url: "https://sheets.googleapis.com/v4/spreadsheets/sheetId/values/A1",
+			} as never),
+		).toBe("https://docs.google.com/spreadsheets/d/sheetId/edit")
+		expect(
+			getDocumentSourceUrl({
+				type: "google_slide",
+				url: "https://slides.googleapis.com/v1/presentations/slideId/pages",
+			} as never),
+		).toBe("https://docs.google.com/presentation/d/slideId/edit")
+	})
+
 	it("maps documents by all known ids", () => {
 		const mapped = mapDocumentsByKnownIds([
 			{ id: "docA", customId: "customA", type: "text", url: null } as never,

apps/web/lib/chat-memory-tools.ts

@@ -206,7 +206,7 @@ export function isMemoryToolOutputReady(
 }
 
 export function extractMemoryToolOutputs(message: {
-	parts?: unknown[]
+	parts?: readonly unknown[]
 }): MemoryToolOutput[] {
 	const parts = Array.isArray(message.parts) ? message.parts : []
 	const outputs: MemoryToolOutput[] = []
@@ -349,11 +349,27 @@ export function getDocumentSourceUrl(
 		customId?: string | null
 	},
 ) {
-	if (document.type === "google_doc" && document.customId)
-		return `https://docs.google.com/document/d/${document.customId}`
-	if (document.type === "google_sheet" && document.customId)
-		return `https://docs.google.com/spreadsheets/d/${document.customId}`
-	if (document.type === "google_slide" && document.customId)
-		return `https://docs.google.com/presentation/d/${document.customId}`
-	return document.url
+	const url = document.url ?? null
+	const googleDocTypes: Record<string, { prefix: string; apiPattern: RegExp }> =
+		{
+			google_doc: {
+				prefix: "https://docs.google.com/document/d/",
+				apiPattern: /docs\.googleapis\.com\/v1\/documents\/([A-Za-z0-9_-]+)/,
+			},
+			google_sheet: {
+				prefix: "https://docs.google.com/spreadsheets/d/",
+				apiPattern:
+					/sheets\.googleapis\.com\/v4\/spreadsheets\/([A-Za-z0-9_-]+)/,
+			},
+			google_slide: {
+				prefix: "https://docs.google.com/presentation/d/",
+				apiPattern:
+					/slides\.googleapis\.com\/v1\/presentations\/([A-Za-z0-9_-]+)/,
+			},
+		}
+	const googleDoc = document.type ? googleDocTypes[document.type] : undefined
+	if (!googleDoc) return url
+	if (document.customId) return `${googleDoc.prefix}${document.customId}/edit`
+	const apiId = url?.match(googleDoc.apiPattern)?.[1]
+	return apiId ? `${googleDoc.prefix}${apiId}/edit` : url
 }