Fix memory tool highlight compatibility

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

@@ -45,11 +45,28 @@ const assistantMessage = {
 } as const
 
 describe("chat memory tool citation mapping", () => {
-	it("extracts only completed memory tool outputs", () => {
-		const outputs = extractMemoryToolOutputs(assistantMessage)
+	it("extracts only ready memory tool outputs", () => {
+		const outputs = extractMemoryToolOutputs({
+			parts: [
+				...assistantMessage.parts,
+				{
+					type: "tool-searchMemories",
+					state: "done",
+					output: { sourceIds: ["done"], documentIds: ["doneDoc"] },
+				},
+				{
+					type: "tool-searchMemories",
+					output: { sourceIds: ["stateless"], documentIds: ["statelessDoc"] },
+				},
+			],
+		})
 
-		expect(outputs).toHaveLength(1)
-		expect(outputs[0]?.output.sourceIds).toEqual(["S1"])
+		expect(outputs).toHaveLength(3)
+		expect(outputs.map((output) => output.output.sourceIds?.[0])).toEqual([
+			"S1",
+			"done",
+			"stateless",
+		])
 	})
 
 	it("maps citation ids to document and custom ids", () => {
@@ -72,6 +89,30 @@ describe("chat memory tool citation mapping", () => {
 		)
 	})
 
+	it("keeps graph highlights for legacy memory tool states and ids", () => {
+		const legacyMessage = {
+			id: "legacy",
+			role: "assistant",
+			parts: [
+				{
+					type: "tool-searchMemories",
+					state: "done",
+					output: { results: [{ id: "legacyDoc" }] },
+				},
+				{
+					type: "tool-recallContext",
+					output: { documentIds: ["statelessDoc"] },
+				},
+			],
+		} as const
+
+		expect(extractMemoryToolOutputs(legacyMessage)).toHaveLength(2)
+		expect(extractHighlightDocumentIdsFromMessages([legacyMessage])).toEqual([
+			"legacyDoc",
+			"statelessDoc",
+		])
+	})
+
 	it("normalizes nested discoverSpaces memory results", () => {
 		const outputs = extractMemoryToolOutputs({
 			parts: [

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

@@ -195,6 +195,16 @@ function firstDocumentForResult(
 	return firstId ? { id: firstId, customId: result.customId } : null
 }
 
+export function isMemoryToolOutputReady(
+	part: Record<string, unknown>,
+): boolean {
+	return (
+		part.state === "output-available" ||
+		part.state === "done" ||
+		(part.state === undefined && part.output !== undefined)
+	)
+}
+
 export function extractMemoryToolOutputs(message: {
 	parts?: unknown[]
 }): MemoryToolOutput[] {
@@ -212,7 +222,7 @@ export function extractMemoryToolOutputs(message: {
 			)
 		)
 			continue
-		if (part.state !== "output-available") continue
+		if (!isMemoryToolOutputReady(part)) continue
 		outputs.push({ output: normalizeOutput(part.output) })
 	}
 
@@ -257,6 +267,7 @@ export function extractDocumentIdsFromMemoryOutput(
 	const ids: string[] = []
 	ids.push(...(output.documentIds ?? []))
 	for (const result of output.results ?? []) {
+		if (result.id) ids.push(result.id)
 		if (result.documentId) ids.push(result.documentId)
 		if (result.internalDocumentId) ids.push(result.internalDocumentId)
 		ids.push(...(result.documentIds ?? []))