bug fixes, and performance improvements, and drawer

tests/citations.test.ts

@@ -0,0 +1,84 @@
+import { describe, expect, test } from "bun:test";
+import { collectCitationSources } from "../app/utils/citations";
+import { createMarkdownParser, getNodeChildren, getNodeTag } from "../app/utils/markdown";
+import type { MarkdownAstNode } from "../app/utils/markdown";
+import { readSearchResponse } from "../server/utils/search";
+
+const citations = (nodes: MarkdownAstNode[]): MarkdownAstNode[] => nodes.flatMap(node =>
+    getNodeTag(node) === "source-cite" ? [node] : citations(getNodeChildren(node)),
+);
+
+describe("citation parsing", () => {
+    test("keeps citation markers as text unless explicitly enabled", async () => {
+        const content = "Claim [[cite:1234]] and [[cite:s7,s9]].";
+        for (const options of [{}, { citations: false }]) {
+            for (const streaming of [false, true]) {
+                const tree = await createMarkdownParser(options)(content, { streaming });
+                expect(citations(tree.nodes)).toEqual([]);
+                expect(getNodeChildren(tree.nodes[0] ?? "")).toEqual([content]);
+            }
+        }
+    });
+
+    test("renders single and grouped references as dedicated nodes", async () => {
+        const tree = await createMarkdownParser({ citations: true })("Claim [[cite:s7]]. Another [[cite:s7,s9]].");
+        expect(citations(tree.nodes)).toEqual([
+            ["source-cite", { ids: "s7" }],
+            ["source-cite", { ids: "s7,s9" }],
+        ]);
+    });
+
+    test("does not interpret code, escaped references, or link labels", async () => {
+        for (const content of [
+            "`[[cite:s7]]`",
+            "```text\n[[cite:s7]]\n```",
+            "\\[[cite:s7]]",
+            "[Claim [[cite:s7]]](https://example.com)",
+        ]) {
+            const tree = await createMarkdownParser({ citations: true })(content);
+            expect(citations(tree.nodes)).toEqual([]);
+        }
+    });
+
+    test("streaming never auto-completes a citation", async () => {
+        const parse = createMarkdownParser({ citations: true });
+        const content = "Claim [[cite:s7,s9]]";
+        for (let length = 1; length < content.length; length++) {
+            expect(citations((await parse(content.slice(0, length), { streaming: true })).nodes)).toEqual([]);
+        }
+        expect(citations((await parse(content, { streaming: true })).nodes)).toHaveLength(1);
+        expect(citations((await parse(content)).nodes)).toHaveLength(1);
+    });
+});
+
+describe("citation sources", () => {
+    test("IDs survive storage, including duplicate URLs and multiple searches", () => {
+        const first = readSearchResponse({ results: [
+            { title: "One", url: "https://example.com" },
+            { title: "Two", url: "https://example.com" },
+        ] });
+        const second = readSearchResponse({ results: [{ title: "Three", url: "https://example.com" }] });
+        expect(first.ok && second.ok).toBe(true);
+        if (!first.ok || !second.ok) {
+            return;
+        }
+        const parts = [first.data, second.data].map(output => ({
+            toolCall: { toolName: "search", output: { type: "json", value: JSON.stringify(output) } },
+        }));
+        const sources = collectCitationSources(parts);
+        expect(sources.size).toBe(3);
+        expect(sources.get(first.data.results[0]?.sourceId ?? "")?.title).toBe("One");
+        expect(collectCitationSources(JSON.parse(JSON.stringify(parts)))).toEqual(sources);
+        expect(sources.has("invented")).toBe(false);
+        expect(collectCitationSources([]).size).toBe(0);
+    });
+
+    test("ignores legacy results, non-search tools, and unsafe URLs", () => {
+        const output = { type: "json", value: JSON.stringify({ results: [
+            { title: "Legacy", link: "https://example.com" },
+            { sourceId: "bad", title: "Bad", link: "javascript:alert(1)" },
+        ] }) };
+        expect(collectCitationSources([{ toolCall: { toolName: "search", output } }]).size).toBe(0);
+        expect(collectCitationSources([{ toolCall: { toolName: "python", output } }]).size).toBe(0);
+    });
+});