pediatric-ai-scribe-v3/test/assistant-citations.test.js

const { test } = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs/promises');
const path = require('node:path');
const MarkdownIt = require('markdown-it');

let modulePromise;

async function loadCitationModule() {
  if (!modulePromise) {
    modulePromise = fs.readFile(path.join(__dirname, '..', 'public', 'js', 'assistant', 'citations.js'), 'utf8')
      .then((source) => import('data:text/javascript;charset=utf-8,' + encodeURIComponent(source)));
  }
  return modulePromise;
}

const sources = [
  { title: 'Nelson Textbook of Pediatrics' },
  { title: 'Pediatric Asthma Guideline' },
  { resource: 'Emergency Medicine Review' }
];

test('renders citation clusters as links to matching source cards', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Give magnesium for severe exacerbation [1, 2].', sources);
  assert.match(html, /href="#assistant-source-1"/);
  assert.match(html, /href="#assistant-source-2"/);
  assert.match(html, /title="Source 1: Nelson Textbook of Pediatrics"/);
  assert.match(html, /data-source-number="1"/);
  assert.match(html, /data-source-number="2"/);
  assert.match(html, /<a class="assistant-cite"[^>]*>src<\/a> <a class="assistant-cite"[^>]*>src<\/a>/);
});

test('renders escaped citation tokens as source links, not display math', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const katex = { renderToString: function () { throw new Error('citation sent to KaTeX'); } };
  const html = renderAssistantMarkdown('Acquired hypothyroidism is uncommon. \\[1, 2\\].', sources, { katex });
  assert.match(html, /<a class="assistant-cite"[^>]*>src<\/a> <a class="assistant-cite"[^>]*>src<\/a>/);
  assert.doesNotMatch(html, /katex-display/);
});

test('can render citation labels as numbers for PDF export', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Give magnesium for severe exacerbation [1, 2].', sources, { citationLabel: 'number' });
  assert.match(html, /<a class="assistant-cite"[^>]*>1<\/a> <a class="assistant-cite"[^>]*>2<\/a>/);
});

test('strips [src]/[source] placeholder tokens instead of inventing numbers', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Give fluids [src] and review in 24h [source].', []);
  assert.doesNotMatch(html, /\[src\]|\[source\]/i);
  assert.match(html, /Give fluids/);
  const kept = renderAssistantMarkdown('Give fluids [1] and review.', [{ number: 1, title: 'S', page: 2 }]);
  assert.match(kept, /data-source-number="1"/, 'real citation numbers still render');
});

test('leaves unknown citation clusters untouched rather than guessing', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Dose statement [4].', sources);
  assert.match(html, /\[4\]/);
  assert.doesNotMatch(html, /assistant-source-4/);
});

test('does not merge separate citations across separate claims', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Use oxygen [1]. Give bronchodilator [2].', sources);
  assert.match(html, /oxygen <a class="assistant-cite"[^>]*data-source-number="1"[^>]*>src<\/a>/);
  assert.match(html, /bronchodilator <a class="assistant-cite"[^>]*data-source-number="2"[^>]*>src<\/a>/);
  assert.doesNotMatch(html, /data-source-number="1"[^>]*>src<\/a><a class="assistant-cite"[^>]*data-source-number="2"/);
});

test('sorts adjacent citation tokens into one safe Vancouver cluster', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Deteriorating course [1][4][2][3].', [
    { title: 'A' }, { title: 'B' }, { title: 'C' }, { title: 'D' }
  ]);
  assert.match(html, /data-source-number="1"[^>]*>src<\/a> <a class="assistant-cite"[^>]*data-source-number="2"[^>]*>src<\/a> <a class="assistant-cite"[^>]*data-source-number="3"[^>]*>src<\/a> <a class="assistant-cite"[^>]*data-source-number="4"[^>]*>src<\/a>/);
  assert.doesNotMatch(html, /\]<span|\]\[/);
});

test('repairs a clearly adjacent trailing citation missing its closing bracket', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Deteriorating course [1][4][2][3', [
    { title: 'A' }, { title: 'B' }, { title: 'C' }, { title: 'D' }
  ]);
  assert.match(html, /data-source-number="1"[^>]*>src<\/a> <a class="assistant-cite"[^>]*data-source-number="2"[^>]*>src<\/a> <a class="assistant-cite"[^>]*data-source-number="3"[^>]*>src<\/a> <a class="assistant-cite"[^>]*data-source-number="4"[^>]*>src<\/a>/);
});

test('does not normalize adjacent citations if any source number is unknown', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Unsupported claim [1][9].', sources);
  assert.match(html, /\[9\]/);
  assert.doesNotMatch(html, /assistant-source-9/);
});

test('escapes source titles inside citation link titles', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Unsafe title [1].', [{ title: 'Bad "Title" <script>' }]);
  assert.match(html, /title="Source 1: Bad &quot;Title&quot; &lt;script&gt;"/);
  assert.doesNotMatch(html, /<script>/);
});

test('does not strip sentences that mention available sources', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Pneumonia can cause parapneumonic effusion [1]. However, the available sources do not specifically identify adenovirus as a typical cause [2]. In summary, evidence is insufficient [3].', sources);
  assert.match(html, /available sources do not specifically identify adenovirus/);
  assert.match(html, /In summary, evidence is insufficient/);
});

test('normalizes inline markdown headings before rendering', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Intro ### Management\n- oxygen', sources);
  assert.match(html, /<h3>Management<\/h3>/);
  assert.match(html, /<li>oxygen<\/li>/);
});

test('normalizes old saved assistant answer formatting', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Comparison of Bronchiolitis and Asthma Management in Infants ### Bronchiolitis Management - Mainstay: Supportive care [1, 2]. - Medications: bronchodilators are not routine [2]. ### Key Differences | Aspect | Bronchiolitis | Asthma | |---|---|---| | Therapy | Supportive | SABA | --- References: [1] duplicate', sources);
  assert.match(html, /<h3>Bronchiolitis Management<\/h3>/);
  assert.match(html, /<li>Mainstay: Supportive care/);
  assert.match(html, /Key Differences/);
});

test('does not turn citation-delimited prose into a list', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Admission is indicated for: Apnoea [1, 2]- **Persistent oxygen saturation <90%** [1, 2].', sources);
  assert.match(html, /<\/a>- <strong>Persistent oxygen saturation &lt;90%<\/strong>/);
  assert.doesNotMatch(html, /<ul>/);
});

test('joins a citation-only paragraph back to its claim', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Acquired hypothyroidism is uncommon.\n\n[1, 2]\n\n.\n\n- Next point', sources);
  assert.match(html, /uncommon\. <a class="assistant-cite"[^>]*>src<\/a> <a class="assistant-cite"[^>]*>src<\/a>\./);
  assert.doesNotMatch(html, /<p>\s*<a class="assistant-cite"/);
});

test('keeps citations inline before clinical terms', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Assess for lethargy [1] and dehydration [2].', sources);
  assert.match(html, /lethargy <a class="assistant-cite"[^>]*>src<\/a> and dehydration <a class="assistant-cite"[^>]*>src<\/a>/);
  assert.doesNotMatch(html, /<\/a><br>\s*and dehydration/);
});

test('keeps citation-delimited clinical prose inline', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('CT is indicated for: [1, 2, 3]- Focal neurologic deficits [1].', sources);
  assert.match(html, /<\/a>- Focal neurologic deficits/);
  assert.doesNotMatch(html, /<br>|<ul>/);
});

test('preserves code block contents without creating citation links inside code', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('```js\nconst ref = "[1]";\n```', sources);
  assert.match(html, /<pre><code>const ref = &quot;\[1\]&quot;/);
  assert.doesNotMatch(html, /assistant-source-1/);
});

test('does not repair source tables inside code blocks', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('```md\n| Feature | Source(s) |\n|---|---|\n| Example | 1, 2 |\n```', sources);
  assert.match(html, /\| Example \| 1, 2 \|/);
  assert.doesNotMatch(html, /assistant-source-1/);
});

test('does not repair model-split markdown tables', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Summary Table\n| Indication for Use | Dose (IV) | Max Dose\n\n| Infusion Time | Monitoring/Precautions |\n|-----------------------------------------|----------------------------|----------|--------------|---------------------------------------|\n| Severe exacerbation unresponsive to SABA/anticholinergic | 25-75 mg/kg | 2 g\n| 20 min | BP, reflexes, respiratory status[1, 2] |', sources);
  assert.doesNotMatch(html, /Max Dose Infusion Time/);
  assert.match(html, /Infusion Time/);
  assert.match(html, /assistant-cite/);
});

test('breaks inline numbered clinical sections onto new lines', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Monitor with serial abdominal exams [1]. 2. Imaging\nImmediate abdominal X-ray for obstruction [1, 2]. 3. Laboratory Studies\nAssess electrolytes if ill-appearing [3]. 4. Surgical Consultation\nUrgent pediatric surgical consultation [1].', sources);
  assert.match(html, /2\. Imaging/);
  assert.match(html, /3\. Laboratory Studies/);
  assert.match(html, /4\. Surgical Consultation/);
  assert.doesNotMatch(html, /exams[^<]*2\. Imaging/);
  assert.doesNotMatch(html, /obstruction[^<]*3\. Laboratory Studies/);
});

test('strips orphan trailing markdown emphasis markers', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const html = renderAssistantMarkdown('Summary Table\n\nAge Group\tEmpiric Antibiotics\n<1 month\tAmpicillin + cefotaxime\n**', sources);
  assert.match(html, /Summary Table/);
  assert.doesNotMatch(html, /\*\*/);
});

test('strips orphan trailing markdown marker after tab-separated text', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const markdownIt = new MarkdownIt({ html: false, linkify: true, typographer: true, breaks: true });
  const input = 'Summary Table\nStep\tRecommendation\nBlood culture\tObtain before antibiotics\nHospitalization\tYes, for all febrile neonates <=28 days old\n**';
  const html = renderAssistantMarkdown(input, sources, { markdownIt });
  assert.doesNotMatch(html, /<table>/);
  assert.match(html, /Hospitalization/);
  assert.doesNotMatch(html, /\*\*/);
});

test('keeps summary paragraph outside tab-separated text', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const markdownIt = new MarkdownIt({ html: false, linkify: true, typographer: true, breaks: true });
  const input = 'Summary Table\nStep/Indication\tRecommendation/Threshold\tCitation\nInitial assessment\tABCs and focused exam\t[1]\nActivated charcoal window\tWithin 1 hour\t[2]\nIn summary: rapidly assess ABCs and consider activated charcoal when appropriate [1, 2].';
  const html = renderAssistantMarkdown(input, sources, { markdownIt });
  assert.doesNotMatch(html, /<table>/);
  assert.match(html, /In summary:/);
  assert.doesNotMatch(html, /<td>In summary:/);
});

test('strips trailing marker after notes followed by tab-separated text', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const markdownIt = new MarkdownIt({ html: false, linkify: true, typographer: true, breaks: true });
  const input = 'Additional Notes\nListeria coverage: Ampicillin should be added for infants <3 months [1].\nVancomycin: Added for resistant pneumococcus [2].\nAge Group\tEmpiric Antibiotics\n<1 month\tAmpicillin + Cefotaxime\n1-23 months\tVancomycin + Cefotaxime or Ceftriaxone\n>=24 months\tVancomycin + Cefotaxime or Ceftriaxone\n**';
  const html = renderAssistantMarkdown(input, sources, { markdownIt });
  assert.match(html, /Additional Notes/);
  assert.doesNotMatch(html, /<table>/);
  assert.doesNotMatch(html, /\*\*/);
});

test('links citation-only cells in markdown pipe tables', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const markdownIt = new MarkdownIt({ html: false, linkify: true, typographer: true, breaks: true });
  const input = '| Step | Purpose | Source(s) |\n|---|---|---|\n| Surgical consult | Urgent assessment | [1, 2] |\n| X-ray | Assess obstruction | [3] |';
  const html = renderAssistantMarkdown(input, sources, { markdownIt });
  assert.match(html, /<table>/);
  assert.match(html, /data-source-number="1"/);
  assert.match(html, /data-source-number="2"/);
  assert.match(html, /data-source-number="3"/);
  assert.doesNotMatch(html, /<td>\[1, 2\]<\/td>/);
});

test('repairs bare source numbers in markdown table source columns', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const markdownIt = new MarkdownIt({ html: false, linkify: true, typographer: true, breaks: true });
  const input = '| Feature | Possible Surgical Cause(s) | Source(s) |\n|---|---|---|\n| Bilious vomiting | Malrotation with volvulus | 1, 2, 3 |\n| Projectile vomiting | Pyloric stenosis | 2 |';
  const html = renderAssistantMarkdown(input, sources, { markdownIt });
  assert.match(html, /<table>/);
  assert.match(html, /data-source-number="1"/);
  assert.match(html, /data-source-number="2"/);
  assert.match(html, /data-source-number="3"/);
  assert.doesNotMatch(html, /<td>1, 2, 3<\/td>/);
});

test('uses markdown-it stack for GFM-style tables when provided', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const markdownIt = new MarkdownIt({ html: false, linkify: true, typographer: true, breaks: true });
  const html = renderAssistantMarkdown('| Age | Antibiotics |\n|---|---|\n| <1 month | Ampicillin + cefotaxime [1] |', sources, { markdownIt });
  assert.match(html, /<table>/);
  assert.match(html, /<td>&lt;1 month<\/td>/);
  assert.match(html, /assistant-cite/);
});

test('renders pipe summary tables and removes trailing emphasis junk', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const markdownIt = new MarkdownIt({ html: false, linkify: true, typographer: true, breaks: true });
  const input = 'Summary Table\n\n| Situation | Imaging Recommended | Timing/Notes |\n|---|---|---|\n| First febrile UTI | Renal/bladder ultrasound | After recovery |\n| Routine use of DMSA scan | Not recommended | |\n**';
  const html = renderAssistantMarkdown(input, sources, { markdownIt });
  assert.match(html, /<table>/);
  assert.match(html, /<th>Situation<\/th>/);
  assert.match(html, /<td>Routine use of DMSA scan<\/td>/);
  assert.doesNotMatch(html, /\*\*/);
});

test('does not repair malformed tab-separated clinical summary rows into a table', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const markdownIt = new MarkdownIt({ html: false, linkify: true, typographer: true, breaks: true });
  const input = 'Summary Table\nStep\tDetails\tSource\nRed Flags\tShock, peritonitis, distension\t3 4 7\nInitial Imaging\tAbdominal X-ray if perforation suspected, then ultrasound\n2 3 7\t\tDiagnostic Gold Std\nAbdominal ultrasound target sign\t2 3 7 Nonoperative Reduction\tAir or contrast enema if stable\nSurgical Indications\tShock, peritonitis, perforation, failed reduction\n1 3 7\t\tRecurrence Rate\n~10% after nonsurgical reduction\t5 6 7\tIf you need details on recurrence management, please specify.';
  const html = renderAssistantMarkdown(input, sources, { markdownIt, citationLabel: 'number' });
  assert.doesNotMatch(html, /<table>/);
  assert.match(html, /Initial Imaging/);
  assert.match(html, /If you need details/);
});

test('clinical assistant streams long table answers as lightweight text before final render', async () => {
  const source = await fs.readFile(path.join(__dirname, '..', 'public', 'js', 'clinicalAssistant.js'), 'utf8');
  assert.match(source, /STREAM_MARKDOWN_LIMIT/);
  assert.match(source, /renderStreamingAnswerHtml\(partial, streamSources\)/);
  assert.match(source, /assistant-streaming-text/);
  assert.match(source, /pipeRows >= 8/);
});

test('mermaid source survives the sanitiser and round-trips', async () => {
  const { renderAssistantMarkdown } = await loadCitationModule();
  const flowchart = 'graph TD; A[Start]-->B[Give O2];';
  const html = renderAssistantMarkdown('```mermaid\n' + flowchart + '\n```', []);
  const attr = html.match(/data-mermaid="([^"]*)"/);
  assert.ok(attr, 'the placeholder must carry the diagram source');
  // Encoded, because DOMPurify strips an attribute whose value contains "-->"
  // and every flowchart has one, which left diagrams stuck on their placeholder.
  assert.ok(!attr[1].includes('-->'), 'the stored value must not contain a raw arrow');
  assert.equal(decodeURIComponent(attr[1].replace(/&amp;/g, '&')).trim(), flowchart);
});

// ── Admin citation mode ────────────────────────────────────────────────────
test('citation mode changes only the citation rules, never the grounding', () => {
  const { buildSystemPrompt } = require('../src/utils/clinicalAnswer');
  const on = buildSystemPrompt('BEHAVIOR');
  const off = buildSystemPrompt('BEHAVIOR', { citations: false });

  assert.match(on, /Cite factual claims immediately with numbered citations/);
  assert.doesNotMatch(off, /Cite factual claims immediately/, 'the [n] rules are dropped');
  assert.match(off, /Do not include citations, source numbers, bracketed reference markers/);
  assert.match(off, /the requirement is unchanged, only the visible marker is dropped/);

  // Everything that makes the answer trustworthy must be identical in both.
  for (const shared of [
    /Answer only the user question/,
    /not directly supported by retrieved sources/,
    /do not claim that another source is from the named source/,
    /valid GitHub-flavored markdown table/,
    /Do not add generic disclaimers/
  ]) {
    assert.match(on, shared);
    assert.match(off, shared, 'grounding and formatting rules are the same in both modes');
  }
  assert.equal(buildSystemPrompt('BEHAVIOR', {}), on, 'citations default to on');
  assert.equal(buildSystemPrompt('BEHAVIOR', { citations: true }), on);
});

test('stray citation markers are stripped from the stored answer, not just the view', () => {
  const { stripCitationMarkers } = require('../src/utils/clinicalAnswer');
  assert.equal(stripCitationMarkers('Amoxicillin 90 mg/kg/day [1][2]. Reassess in 48 h [3].'),
    'Amoxicillin 90 mg/kg/day. Reassess in 48 h.');
  assert.equal(stripCitationMarkers('Give fluids [1, 3] and rest [2].'), 'Give fluids and rest.');
  assert.equal(stripCitationMarkers('No citations here.'), 'No citations here.');
  // A model that obeys the prompt costs nothing; one that does not must not
  // leave literal "[1]" text with no sources panel to click through to.
  assert.doesNotMatch(stripCitationMarkers('Dose [1] is 90 mg/kg [2]'), /\[\d/);
});

test('the admin toggle is a real boolean setting and the server refuses anything else', () => {
  const fs = require('node:fs');
  const path = require('node:path');
  const root = path.join(__dirname, '..');
  const admin = fs.readFileSync(path.join(root, 'src/routes/adminConfig.js'), 'utf8');
  assert.match(admin, /clinical_assistant\.citations_enabled' && !\['true', 'false'\]/, 'validated as a boolean');
  const route = fs.readFileSync(path.join(root, 'src/routes/clinicalAssistant.js'), 'utf8');
  assert.match(route, /citations_enabled', 'true'\)\) !== 'false'/, 'defaults to on');
  assert.match(route, /citationsEnabled: await citationsEnabled\(\)/, 'exposed on the status endpoint');
  assert.match(route, /sources: prepared\.citations \? sanitizeSourcesForClient\(prepared\.sources\) : \[\]/,
    'no sources are sent to the client when citations are off');
  const html = fs.readFileSync(path.join(root, 'public/components/admin.html'), 'utf8');
  assert.match(html, /id="assistant-citations-enabled"/, 'the admin has a control for it');
  const ui = fs.readFileSync(path.join(root, 'public/js/clinicalAssistant.js'), 'utf8');
  assert.match(ui, /applyCitationMode\(data\.citationsEnabled !== false\)/, 'the UI follows the server');
  assert.match(ui, /side\.hidden = !citationsOn/, 'and hides the Sources panel');
});