diff --git a/apps/web/components/chat/index.tsx b/apps/web/components/chat/index.tsx index a2d8e305..aae5920b 100644 --- a/apps/web/components/chat/index.tsx +++ b/apps/web/components/chat/index.tsx @@ -48,7 +48,10 @@ import { import { SpaceSelector } from "@/components/space-selector" import { SuperLoader } from "../superloader" import { UserMessage } from "./message/user-message" -import { AgentMessage } from "./message/agent-message" +import { + AgentMessage, + isChatToolDisplayPartType, +} from "./message/agent-message" import { ChatGraphContextRail } from "./chat-graph-context-rail" import { ChainOfThought } from "./input/chain-of-thought" import { useIsMobile } from "@hooks/use-mobile" @@ -1147,14 +1150,14 @@ export function ChatSidebar({ }) => ({ id: m.id, role: m.role, - // Strip tool parts (they break convertToModelMessages with tool_use/tool_result - // mismatches); keep text/reasoning + source parts so citations survive reload. + // Keep chat tool outputs that are meaningful to render after thread reload. parts: (m.parts || []).filter( (p) => p.type === "text" || p.type === "reasoning" || p.type === "source-url" || - p.type === "source-document", + p.type === "source-document" || + isChatToolDisplayPartType(p.type), ), metadata: m.metadata, createdAt: new Date(m.createdAt), diff --git a/apps/web/components/chat/message/agent-message.tsx b/apps/web/components/chat/message/agent-message.tsx index 1b7f69a8..eac0f7b9 100644 --- a/apps/web/components/chat/message/agent-message.tsx +++ b/apps/web/components/chat/message/agent-message.tsx @@ -2,6 +2,7 @@ import { type ReactNode, useEffect, useMemo, useRef, useState } from "react" import type { UIMessage } from "@ai-sdk/react" +import { useQuery } from "@tanstack/react-query" import { Streamdown } from "streamdown" import { BookOpenIcon, @@ -19,18 +20,37 @@ import { } from "lucide-react" import { cn } from "@lib/utils" import { isWebSearchToolName } from "@/lib/chat-web-search-tools" +import { + buildCitationIndex, + fetchDocumentsByIds, + getDocumentSourceUrl, + isMemoryToolOutputReady, + mapDocumentsByKnownIds, + type CitationTarget, + type DocumentWithMemories, + extractMemoryToolOutputs, +} from "@/lib/chat-memory-tools" +import { + parseSourceAnnotatedMarkdown, + stripSourceMarkup, +} from "@/lib/source-annotations" import { modelNames, type ModelId } from "@/lib/models" import { RelatedMemories } from "./related-memories" import { MessageActions } from "./message-actions" const TOOL_META: Record = { bash: { label: "Memory", icon: TerminalIcon }, + recallContext: { label: "Recall Memories", icon: BookOpenIcon }, + discoverSpaces: { label: "Discover Spaces", icon: SearchIcon }, web_search: { label: "Web search", icon: GlobeIcon }, google_search: { label: "Google search", icon: GlobeIcon }, // legacy tool names kept for existing persisted messages searchMemories: { label: "Search Memories", icon: SearchIcon }, addMemory: { label: "Add Memory", icon: PlusIcon }, fetchMemory: { label: "Fetch Memory", icon: BookOpenIcon }, + forgetMemory: { label: "Forget Memory", icon: XCircleIcon }, + updateMemory: { label: "Update Memory", icon: BookOpenIcon }, + forgetDocument: { label: "Forget Document", icon: XCircleIcon }, scheduleTask: { label: "Schedule Task", icon: ClockIcon }, listSchedules: { label: "List Schedules", icon: ListIcon }, cancelSchedule: { label: "Cancel Schedule", icon: XCircleIcon }, @@ -38,7 +58,7 @@ const TOOL_META: Record = { type ToolCallDisplayPart = { type: string - state: string + state?: string input?: unknown output?: unknown toolCallId?: string @@ -64,6 +84,20 @@ function faviconUrl(host: string): string { return `https://www.google.com/s2/favicons?sz=64&domain=${host}` } +function safeExternalUrl(url: string | null | undefined): string | null { + if (!url) return null + if (url.startsWith("/") && !url.startsWith("//")) return url + if (url.startsWith("#") && !url.startsWith("#sm-source:")) return url + try { + const parsed = new URL(url) + return parsed.protocol === "http:" || parsed.protocol === "https:" + ? url + : null + } catch { + return null + } +} + function isWebSearchPart(part: { type: string; toolName?: string }): boolean { if (part.type === "dynamic-tool") { return isWebSearchToolName(part.toolName ?? "") @@ -74,6 +108,25 @@ function isWebSearchPart(part: { type: string; toolName?: string }): boolean { return false } +function isMemoryRetrievalToolName(toolName: string): boolean { + return ( + toolName === "searchMemories" || + toolName === "recallContext" || + toolName === "discoverSpaces" + ) +} + +export function isChatToolDisplayPartType(type: string): boolean { + return ( + type === "tool-searchMemories" || + type === "tool-recallContext" || + type === "tool-discoverSpaces" || + type === "tool-forgetMemory" || + type === "tool-updateMemory" || + type === "tool-forgetDocument" + ) +} + function CitationLink({ href, label, @@ -83,7 +136,8 @@ function CitationLink({ label: string source?: SourceUrlPart }) { - const url = source?.url ?? href + const url = safeExternalUrl(source?.url ?? href) ?? "" + if (!url) return <>{label} const host = sourceHost(url) const rawTitle = source?.title?.trim() const hasTitle = @@ -136,9 +190,138 @@ function CitationLink({ ) } -function makeMarkdownComponents(sources: SourceUrlPart[]) { +function sourceTitle( + target: CitationTarget, + document?: DocumentWithMemories, +): string { + return ( + document?.title?.trim() || + target.title?.trim() || + document?.customId || + target.customId || + target.documentId || + target.sourceId + ) +} + +function sourceSummary( + target: CitationTarget, + document?: DocumentWithMemories, +): string | null { + const summary = + document?.summary || + target.summary || + (document as { content?: string } | undefined)?.content || + null + return summary ? summary.trim() : null +} + +function sourceKind( + target: CitationTarget, + document?: DocumentWithMemories, +): string { + return (document?.type || target.type || "memory").replaceAll("_", " ") +} + +function SourceCitationLink({ + sourceId, + children, + citationIndex, + documentByKnownId, +}: { + sourceId: string + children: ReactNode + citationIndex: Map + documentByKnownId: Map +}) { + const target = citationIndex.get(sourceId) + if (!target) return <>{children} + + const document = + (target.documentId + ? documentByKnownId.get(target.documentId) + : undefined) ?? + (target.customId ? documentByKnownId.get(target.customId) : undefined) + const url = safeExternalUrl( + document ? getDocumentSourceUrl(document) : target.url, + ) + const title = sourceTitle(target, document) + const summary = sourceSummary(target, document) + + return ( + + {url ? ( + + {children} + + ) : ( + + )} + + {sourceId} + + + + + + {title} + + + {sourceKind(target, document)} + + + {summary ? ( + + {summary} + + ) : null} + {url ? ( + + Open source + + ) : null} + + + + ) +} + +function makeMarkdownComponents( + sources: SourceUrlPart[], + citationIndex: Map, + documentByKnownId: Map, +) { return { a: ({ href, children }: { href?: string; children?: ReactNode }) => { + if (href?.startsWith("#sm-source:")) { + const sourceId = (() => { + try { + return decodeURIComponent(href.slice("#sm-source:".length)) + } catch { + return null + } + })() + if (!sourceId) return <>{children} + return ( + + {children} + + ) + } const label = typeof children === "string" ? children @@ -146,14 +329,16 @@ function makeMarkdownComponents(sources: SourceUrlPart[]) { ? children.join("") : "" const match = label.match(/^\[?(\d+)\]?$/) - if (match && href) { + const safeHref = safeExternalUrl(href) + if (match && safeHref) { const n = Number(match[1]) - const source = sources.find((s) => s.url === href) ?? sources[n - 1] - return + const source = sources.find((s) => s.url === safeHref) ?? sources[n - 1] + return } + if (!safeHref) return <>{children} return ( @@ -516,7 +704,38 @@ export function AgentMessage({ .filter((part) => part.type === "text") .map((part) => part.text) .join(" ") - const webSources = (() => { + const copyText = stripSourceMarkup(messageText) + const memoryOutputs = useMemo( + () => extractMemoryToolOutputs(message), + [message], + ) + const citationIndex = useMemo( + () => buildCitationIndex(memoryOutputs), + [memoryOutputs], + ) + const allowedSourceIds = useMemo( + () => new Set(citationIndex.keys()), + [citationIndex], + ) + const sourceDocumentIds = useMemo(() => { + const ids = new Set() + for (const target of citationIndex.values()) { + if (target.documentId) ids.add(target.documentId) + if (target.customId) ids.add(target.customId) + } + return [...ids].sort() + }, [citationIndex]) + const { data: sourceDocuments = [] } = useQuery({ + queryKey: ["chat-source-documents", sourceDocumentIds], + queryFn: () => fetchDocumentsByIds(sourceDocumentIds), + enabled: sourceDocumentIds.length > 0, + staleTime: 5 * 60 * 1000, + }) + const documentByKnownId = useMemo( + () => mapDocumentsByKnownIds(sourceDocuments), + [sourceDocuments], + ) + const webSources = useMemo(() => { const seen = new Set() const out: SourceUrlPart[] = [] for (const part of message.parts) { @@ -527,15 +746,13 @@ export function AgentMessage({ out.push(source) } return out - })() + }, [message.parts]) const hasAssistantText = message.parts.some( (p) => p.type === "text" && (p as { text?: string }).text?.trim(), ) - const sourceKey = webSources.map((s) => s.url).join("|") - // biome-ignore lint/correctness/useExhaustiveDependencies: keyed by stable source urls const markdownComponents = useMemo( - () => makeMarkdownComponents(webSources), - [sourceKey], + () => makeMarkdownComponents(webSources, citationIndex, documentByKnownId), + [webSources, citationIndex, documentByKnownId], ) const responseModelLabel = responseModel ? `${modelNames[responseModel].name} ${modelNames[responseModel].version}` @@ -603,7 +820,10 @@ export function AgentMessage({ className="text-sm text-white/90 chat-markdown-content" > - {runText} + { + parseSourceAnnotatedMarkdown(runText, allowedSourceIds) + .markdown + } ) @@ -613,7 +833,7 @@ export function AgentMessage({ type: "dynamic-tool" toolName: string toolCallId: string - state: string + state?: string input?: unknown output?: unknown errorText?: string @@ -651,7 +871,7 @@ export function AgentMessage({
() // [doc:] annotations from sgrep --include-doc-ids (highest confidence) for (const m of text.matchAll(DOC_ANNOTATION)) { - found.add(m[1]) + const id = m[1] + if (id) found.add(id) } // Standard UUID format for (const m of text.matchAll(UUID_IN_STRING)) { @@ -52,7 +57,8 @@ export function documentIdsFromBashText(text: string): string[] { const quoted = /"documentId"\s*:\s*"([^"]+)"/g let q = quoted.exec(text) while (q !== null) { - found.add(q[1]) + const id = q[1] + if (id) found.add(id) q = quoted.exec(text) } return [...found] @@ -136,9 +142,23 @@ export function extractHighlightDocumentIdsFromMessages( if (message.role !== "assistant") continue const parts = message.parts if (!parts) continue + for (const memoryOutput of extractMemoryToolOutputs(message)) { + for (const id of extractDocumentIdsFromMemoryOutput( + memoryOutput.output, + )) { + ids.add(id) + } + } for (const part of parts) { const p = part as Record + if ( + p.type === "tool-searchMemories" || + p.type === "tool-recallContext" || + p.type === "tool-discoverSpaces" + ) { + continue + } if (p.type === "source-document") { const sid = (p as { sourceId?: unknown }).sourceId diff --git a/apps/web/lib/chat-memory-tools.test.ts b/apps/web/lib/chat-memory-tools.test.ts new file mode 100644 index 00000000..9b6411c7 --- /dev/null +++ b/apps/web/lib/chat-memory-tools.test.ts @@ -0,0 +1,177 @@ +import { describe, expect, it } from "bun:test" +import { extractHighlightDocumentIdsFromMessages } from "./chat-highlight-documents" +import { + buildCitationIndex, + extractDocumentIdsFromMemoryOutput, + extractMemoryToolOutputs, + getDocumentSourceUrl, + mapDocumentsByKnownIds, +} from "./chat-memory-tools" + +const assistantMessage = { + id: "m1", + role: "assistant", + parts: [ + { + type: "tool-recallContext", + state: "output-available", + output: { + sourceIds: ["S1"], + documentIds: ["topDoc"], + results: [ + { + citationId: "S1", + content: "memo", + document: { + id: "docA", + customId: "customA", + title: "Doc A", + type: "google_doc", + summary: "sum", + }, + }, + ], + }, + }, + { + type: "tool-discoverSpaces", + state: "input-streaming", + output: { sourceIds: ["ignored"], documentIds: ["ignoredDoc"] }, + }, + { + type: "text", + text: 'Answer from memory', + }, + ], +} as const + +describe("chat memory tool citation mapping", () => { + 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(3) + expect(outputs.map((output) => output.output.sourceIds?.[0])).toEqual([ + "S1", + "done", + "stateless", + ]) + }) + + it("maps citation ids to document and custom ids", () => { + const [output] = extractMemoryToolOutputs(assistantMessage) + const index = buildCitationIndex(output ? [output] : []) + + expect(index.get("S1")?.documentId).toBe("docA") + expect(index.get("S1")?.customId).toBe("customA") + expect(index.has("ignored")).toBe(false) + }) + + it("extracts graph highlight document ids from memory outputs", () => { + const [output] = extractMemoryToolOutputs(assistantMessage) + + expect(output && extractDocumentIdsFromMemoryOutput(output.output)).toEqual( + ["topDoc", "docA", "customA"], + ) + expect( + extractHighlightDocumentIdsFromMessages([assistantMessage as never]), + ).toEqual(["topDoc", "docA", "customA"]) + }) + + 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 as never]), + ).toEqual(["legacyDoc", "statelessDoc"]) + }) + + it("normalizes nested discoverSpaces memory results", () => { + const outputs = extractMemoryToolOutputs({ + parts: [ + { + type: "tool-discoverSpaces", + state: "output-available", + output: { + spaces: [ + { + sourceIds: ["S2"], + documentIds: ["spaceDoc"], + results: [{ citationId: "S2", documentIds: ["nestedDoc"] }], + }, + ], + }, + }, + ], + }) + + const index = buildCitationIndex(outputs) + expect(index.get("S2")?.documentId).toBe("nestedDoc") + expect( + extractDocumentIdsFromMemoryOutput(outputs[0]?.output ?? {}), + ).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, + ]) + expect(mapped.get("docA")?.id).toBe("docA") + expect(mapped.get("customA")?.id).toBe("docA") + }) +}) diff --git a/apps/web/lib/chat-memory-tools.ts b/apps/web/lib/chat-memory-tools.ts new file mode 100644 index 00000000..140c2ae6 --- /dev/null +++ b/apps/web/lib/chat-memory-tools.ts @@ -0,0 +1,375 @@ +import { $fetch } from "@lib/api" +import type { DocumentsWithMemoriesResponseSchema } from "@repo/validation/api" +import type { z } from "zod" +import { isSafeSourceId } from "./source-annotations" + +export const MEMORY_TOOL_PART_TYPES = [ + "tool-searchMemories", + "tool-recallContext", + "tool-discoverSpaces", +] as const +export const MAX_INLINE_GRAPH_DOCUMENT_IDS = 20 + +export type MemoryToolName = + | "searchMemories" + | "recallContext" + | "discoverSpaces" +export type ToolDocumentMetadata = { + id?: string | undefined + internalDocumentId?: string | undefined + customId?: string | null | undefined + title?: string | null | undefined + type?: string | null | undefined + summary?: string | null | undefined + url?: string | null | undefined +} +export type MemoryToolResultItem = { + id?: string | undefined + citationId?: string | undefined + kind?: "memory" | "chunk" | "aggregate" | string | undefined + content?: string | undefined + score?: number | undefined + documentId?: string | undefined + documentIds?: string[] | undefined + internalDocumentId?: string | undefined + customId?: string | undefined + documents?: ToolDocumentMetadata[] | undefined + document?: ToolDocumentMetadata | undefined +} +export type MemoryToolResult = { + query?: string | undefined + count?: number | undefined + sourceIds?: string[] | undefined + documentIds?: string[] | undefined + results?: MemoryToolResultItem[] | undefined + spaces?: + | Array<{ + results?: MemoryToolResultItem[] | undefined + sourceIds?: string[] | undefined + documentIds?: string[] | undefined + }> + | undefined +} +export type MemoryToolOutput = { + output: MemoryToolResult +} +export type CitationTarget = { + sourceId: string + documentId?: string | undefined + customId?: string | null | undefined + title?: string | null | undefined + type?: string | null | undefined + summary?: string | null | undefined + url?: string | null | undefined +} + +export type DocumentWithMemories = z.infer< + typeof DocumentsWithMemoriesResponseSchema +>["documents"][0] + +function isObject(value: unknown): value is Record { + return typeof value === "object" && value !== null +} + +function strings(values: unknown): string[] { + if (!Array.isArray(values)) return [] + return values.filter( + (value): value is string => typeof value === "string" && value.length > 0, + ) +} + +function normalizeOutput(output: unknown): MemoryToolResult { + if (!isObject(output)) return {} + + const nested = [output] + for (const key of [ + "memory", + "search", + "searchResult", + "memoryResult", + "memoryOutput", + "hints", + ]) { + const value = output[key] + if (isObject(value)) nested.push(value) + } + + const sourceIds: string[] = [] + const documentIds: string[] = [] + const results: MemoryToolResultItem[] = [] + + const merged: MemoryToolResult = { + query: typeof output.query === "string" ? output.query : undefined, + count: typeof output.count === "number" ? output.count : undefined, + sourceIds, + documentIds, + results, + spaces: Array.isArray(output.spaces) + ? (output.spaces.filter(isObject) as MemoryToolResult["spaces"]) + : undefined, + } + + for (const value of nested) { + sourceIds.push(...strings(value.sourceIds)) + documentIds.push(...strings(value.documentIds)) + if (Array.isArray(value.results)) + results.push( + ...(value.results.filter(isObject) as MemoryToolResultItem[]), + ) + } + + for (const space of merged.spaces ?? []) { + sourceIds.push(...strings(space.sourceIds)) + documentIds.push(...strings(space.documentIds)) + if (Array.isArray(space.results)) + results.push( + ...(space.results.filter(isObject) as MemoryToolResultItem[]), + ) + } + + merged.sourceIds = dedupe(merged.sourceIds ?? []) + merged.documentIds = dedupe(merged.documentIds ?? []) + return merged +} + +function dedupe(values: string[]): string[] { + return Array.from(new Set(values.filter(Boolean))) +} + +function addTarget( + index: Map, + key: unknown, + target: CitationTarget, +) { + if (typeof key !== "string" || !isSafeSourceId(key)) return + if (!index.has(key)) index.set(key, { ...target, sourceId: key }) +} + +function citationTargetForResult( + sourceId: string, + result: MemoryToolResultItem, +): CitationTarget { + const doc = firstDocumentForResult(result) + const docId = + doc?.internalDocumentId ?? + result.internalDocumentId ?? + doc?.id ?? + result.documentIds?.find(Boolean) ?? + result.documentId + const target = documentTarget(sourceId, doc) + target.documentId = target.documentId ?? docId + target.customId = target.customId ?? result.customId + return target +} + +function documentTarget( + sourceId: string, + doc?: ToolDocumentMetadata | null, +): CitationTarget { + return { + sourceId, + documentId: doc?.internalDocumentId ?? doc?.id, + customId: + doc?.customId ?? + (doc?.internalDocumentId && doc.id !== doc.internalDocumentId + ? doc.id + : undefined), + title: doc?.title, + type: doc?.type, + summary: doc?.summary, + url: doc?.url, + } +} + +function firstDocumentForResult( + result: MemoryToolResultItem, +): ToolDocumentMetadata | null { + if (result.document && isObject(result.document)) return result.document + if (Array.isArray(result.documents) && result.documents.length > 0) + return result.documents.find(isObject) ?? null + const firstId = + result.documentIds?.find(Boolean) ?? + result.documentId ?? + result.internalDocumentId ?? + result.customId + return firstId ? { id: firstId, customId: result.customId } : null +} + +export function isMemoryToolOutputReady( + part: Record, +): boolean { + return ( + part.state === "output-available" || + part.state === "done" || + (part.state === undefined && part.output !== undefined) + ) +} + +export function extractMemoryToolOutputs(message: { + parts?: readonly unknown[] +}): MemoryToolOutput[] { + const parts = Array.isArray(message.parts) ? message.parts : [] + const outputs: MemoryToolOutput[] = [] + + for (let partIndex = 0; partIndex < parts.length; partIndex++) { + const part = parts[partIndex] + if (!isObject(part)) continue + const type = part.type + if ( + typeof type !== "string" || + !MEMORY_TOOL_PART_TYPES.includes( + type as (typeof MEMORY_TOOL_PART_TYPES)[number], + ) + ) + continue + if (!isMemoryToolOutputReady(part)) continue + outputs.push({ output: normalizeOutput(part.output) }) + } + + return outputs +} + +export function buildCitationIndex( + outputs: MemoryToolOutput[], +): Map { + const index = new Map() + + for (const { output } of outputs) { + for (const result of output.results ?? []) { + if (result.citationId) + addTarget( + index, + result.citationId, + citationTargetForResult(result.citationId, result), + ) + } + + for (const sourceId of output.sourceIds ?? []) { + if (index.has(sourceId)) continue + const matchingResult = (output.results ?? []).find( + (result) => result.citationId === sourceId, + ) + if (matchingResult) + addTarget( + index, + sourceId, + citationTargetForResult(sourceId, matchingResult), + ) + } + } + + return index +} + +export function extractDocumentIdsFromMemoryOutput( + output: MemoryToolResult, +): string[] { + 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 ?? [])) + if (result.document?.id) ids.push(result.document.id) + if (result.document?.customId) ids.push(result.document.customId) + for (const doc of result.documents ?? []) { + if (doc.internalDocumentId) ids.push(doc.internalDocumentId) + if (doc.id) ids.push(doc.id) + if (doc.customId) ids.push(doc.customId) + } + } + return dedupe(ids).slice(0, MAX_INLINE_GRAPH_DOCUMENT_IDS) +} + +export async function fetchDocumentsByIds( + ids: string[], +): Promise { + const uniqueIds = dedupe(ids) + if (uniqueIds.length === 0) return [] + + const fetchBy = async (by: "id" | "customId", requestedIds: string[]) => { + const response = await $fetch("@post/documents/documents/by-ids", { + body: { + ids: requestedIds, + by, + }, + disableValidation: true, + }) + const result = response as { + error?: { message?: string } | null + data?: { documents?: DocumentWithMemories[] } | null + } + if (result.error) { + throw new Error("Failed to fetch source documents", { + cause: result.error, + }) + } + return result.data?.documents ?? [] + } + + const byIdDocs = await fetchBy("id", uniqueIds) + const seen = new Set() + const foundLookup = new Set() + for (const doc of byIdDocs) { + if (doc.id) { + seen.add(doc.id) + foundLookup.add(doc.id) + } + if (doc.customId) foundLookup.add(doc.customId) + } + + const unresolved = uniqueIds.filter((id) => !foundLookup.has(id)) + const byCustomDocs = + unresolved.length > 0 ? await fetchBy("customId", unresolved) : [] + const merged = [...byIdDocs] + for (const doc of byCustomDocs) { + if (doc.id && !seen.has(doc.id)) { + seen.add(doc.id) + merged.push(doc) + } + } + + return merged +} + +export function mapDocumentsByKnownIds( + documents: DocumentWithMemories[], +): Map { + const map = new Map() + for (const doc of documents) { + if (doc.id) map.set(doc.id, doc) + if (doc.customId) map.set(doc.customId, doc) + } + return map +} + +export function getDocumentSourceUrl( + document: Pick & { + customId?: string | null + }, +) { + const url = document.url ?? null + const googleDocTypes: Record = + { + 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 +} diff --git a/apps/web/lib/source-annotations.test.ts b/apps/web/lib/source-annotations.test.ts new file mode 100644 index 00000000..dce37e1a --- /dev/null +++ b/apps/web/lib/source-annotations.test.ts @@ -0,0 +1,92 @@ +import { describe, expect, it } from "bun:test" +import { + isSafeSourceId, + parseSourceAnnotatedMarkdown, + stripSourceMarkup, +} from "./source-annotations" + +describe("source annotation parsing", () => { + it("turns allowed response source spans into internal citation links", () => { + const parsed = parseSourceAnnotatedMarkdown( + 'Alpha Beta [x] Gamma', + new Set(["S1"]), + ) + + expect(parsed.markdown).toBe("Alpha [Beta \\[x\\]](#sm-source:S1) Gamma") + }) + + it("renders repeated allowed citations as separate internal links", () => { + const parsed = parseSourceAnnotatedMarkdown( + 'First and Second', + new Set(["S1"]), + ) + + expect(parsed.markdown).toBe( + "[First](#sm-source:S1) and [Second](#sm-source:S1)", + ) + }) + + it("renders unknown or unsafe source ids as plain text", () => { + expect( + parseSourceAnnotatedMarkdown( + 'Unknown plain', + new Set(["S1"]), + ).markdown, + ).toBe("Unknown plain") + + expect( + parseSourceAnnotatedMarkdown( + 'Unsafe plain', + new Set(["bad/id"]), + ).markdown, + ).toBe("Unsafe plain") + }) + + it("keeps unclosed, incomplete, nested, and malformed source markup safe", () => { + expect( + parseSourceAnnotatedMarkdown( + 'Lead unfinished answer', + new Set(["S1"]), + ).markdown, + ).toBe("Lead unfinished answer") + + expect( + parseSourceAnnotatedMarkdown("Lead A B C', + new Set(["S1", "S2"]), + ).markdown, + ).toBe("Outer A B C") + + expect( + parseSourceAnnotatedMarkdown( + "Malformed plain", + new Set(["S1"]), + ).markdown, + ).toBe("Malformed plain") + }) + + it("does not mutate inline code or fenced code blocks", () => { + const input = + '`code`\n```\nfenced\n```' + + expect(parseSourceAnnotatedMarkdown(input, new Set(["S1"])).markdown).toBe( + input, + ) + }) + + it("strips source markup for copy text", () => { + expect( + stripSourceMarkup('Alpha Beta'), + ).toBe("Alpha Beta") + }) + + it("allows only source ids that are safe in internal fragments", () => { + expect(isSafeSourceId("S1._:-")).toBe(true) + expect(isSafeSourceId("bad/id")).toBe(false) + expect(isSafeSourceId("bad space")).toBe(false) + }) +}) diff --git a/apps/web/lib/source-annotations.ts b/apps/web/lib/source-annotations.ts new file mode 100644 index 00000000..7abe6859 --- /dev/null +++ b/apps/web/lib/source-annotations.ts @@ -0,0 +1,194 @@ +export type ParsedSourceAnnotations = { + markdown: string +} + +const RESPONSE_OPEN_PREFIX = " 0 && SAFE_SOURCE_ID_RE.test(id) +} + +function escapeMarkdownLinkText(text: string): string { + return text.replace(/([\\[\]])/g, "\\$1").replace(/\n/g, " ") +} + +function parseOpeningTag( + text: string, + index: number, +): { end: number; sourceId: string } | null | "incomplete" { + if (!text.startsWith(RESPONSE_OPEN_PREFIX, index)) return null + + const tagEnd = text.indexOf(">", index + RESPONSE_OPEN_PREFIX.length) + if (tagEnd === -1) return "incomplete" + + const rawTag = text.slice(index, tagEnd + 1) + const inside = rawTag.slice(1, -1).trim() + if (!inside.startsWith("response")) return null + + let cursor = "response".length + while ( + inside[cursor] === " " || + inside[cursor] === "\t" || + inside[cursor] === "\n" || + inside[cursor] === "\r" + ) + cursor++ + if (!inside.startsWith(SOURCE_ATTR_PREFIX, cursor)) return null + cursor += SOURCE_ATTR_PREFIX.length + + const sourceEnd = inside.indexOf('"', cursor) + if (sourceEnd === -1) return null + const sourceId = inside.slice(cursor, sourceEnd) + cursor = sourceEnd + 1 + while ( + inside[cursor] === " " || + inside[cursor] === "\t" || + inside[cursor] === "\n" || + inside[cursor] === "\r" + ) + cursor++ + if (cursor !== inside.length) return null + if (!isSafeSourceId(sourceId)) return null + + return { end: tagEnd + 1, sourceId } +} + +function advanceCodeState( + text: string, + index: number, + state: { inFence: boolean; inInlineCode: boolean; lineStart: boolean }, +): boolean { + if (state.lineStart && text.startsWith("```", index)) { + state.inFence = !state.inFence + return true + } + + if (!state.inFence && text[index] === "`") { + state.inInlineCode = !state.inInlineCode + return true + } + + return false +} + +function appendChar( + text: string, + index: number, + output: string[], + state: { lineStart: boolean }, +) { + const ch = text[index] ?? "" + output.push(ch) + state.lineStart = ch === "\n" +} + +export function parseSourceAnnotatedMarkdown( + text: string, + allowedSourceIds: ReadonlySet, +): ParsedSourceAnnotations { + const output: string[] = [] + const codeState = { inFence: false, inInlineCode: false, lineStart: true } + + let i = 0 + while (i < text.length) { + if (advanceCodeState(text, i, codeState)) { + appendChar(text, i, output, codeState) + i++ + continue + } + + if ( + !codeState.inFence && + !codeState.inInlineCode && + text.startsWith(RESPONSE_OPEN_PREFIX, i) + ) { + const opening = parseOpeningTag(text, i) + if (opening === "incomplete") { + break + } + + if (opening) { + const closeIndex = text.indexOf(RESPONSE_CLOSE_TAG, opening.end) + if (closeIndex === -1) { + output.push(stripSourceMarkup(text.slice(opening.end))) + break + } + + const inner = text.slice(opening.end, closeIndex) + const hasNested = + inner.includes(RESPONSE_OPEN_PREFIX) || + inner.includes(RESPONSE_CLOSE_TAG) + const isAllowed = allowedSourceIds.has(opening.sourceId) + + if (hasNested) { + const outerCloseIndex = text.indexOf( + RESPONSE_CLOSE_TAG, + closeIndex + RESPONSE_CLOSE_TAG.length, + ) + const fallbackEnd = + outerCloseIndex === -1 ? closeIndex : outerCloseIndex + output.push(stripSourceMarkup(text.slice(opening.end, fallbackEnd))) + i = fallbackEnd + RESPONSE_CLOSE_TAG.length + continue + } + + const plainInner = stripSourceMarkup(inner) + if (isAllowed && plainInner.trim().length > 0) { + output.push( + `[${escapeMarkdownLinkText(plainInner)}](#sm-source:${encodeURIComponent(opening.sourceId)})`, + ) + } else { + output.push(plainInner) + } + + i = closeIndex + RESPONSE_CLOSE_TAG.length + codeState.lineStart = + output.length === 0 || + output[output.length - 1]?.endsWith("\n") === true + continue + } + + const nextClose = text.indexOf(RESPONSE_CLOSE_TAG, i) + if (nextClose !== -1) { + const tagEnd = text.indexOf(">", i) + if (tagEnd !== -1 && tagEnd < nextClose) { + output.push(stripSourceMarkup(text.slice(tagEnd + 1, nextClose))) + i = nextClose + RESPONSE_CLOSE_TAG.length + continue + } + } + } + + appendChar(text, i, output, codeState) + i++ + } + + return { markdown: output.join("") } +} + +export function stripSourceMarkup(text: string): string { + let output = "" + let i = 0 + + while (i < text.length) { + if (text.startsWith(RESPONSE_CLOSE_TAG, i)) { + i += RESPONSE_CLOSE_TAG.length + continue + } + + if (text.startsWith(RESPONSE_OPEN_PREFIX, i)) { + const tagEnd = text.indexOf(">", i + RESPONSE_OPEN_PREFIX.length) + if (tagEnd === -1) break + i = tagEnd + 1 + continue + } + + output += text[i] + i++ + } + + return output +}