pediatric-ai-scribe-v3/test/assistant-attachment-roundtrip.test.js

const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const path = require('node:path');
const vm = require('node:vm');
const express = require('express');
const { JSDOM } = require('jsdom');
const { marked } = require('marked');
const policy = require('../src/utils/clinicalConversation');
const answer = require('../src/utils/clinicalAnswer');
const root = path.join(__dirname, '..');
const read = file => fs.readFileSync(path.join(root, file), 'utf8');
const quiet = { log() {}, warn() {}, error() {}, info() {} };
const canonical = (length, fill = 7) => Buffer.alloc(length, fill).toString('base64');
const png = { dataBase64: canonical(16), mimeType: 'image/png', name: 'chest.png' };
const jpeg = { dataBase64: canonical(32), mimeType: 'image/jpeg', name: 'knee.jpeg' };
const asset = '/api/generated-images/12345678-1234-1234-1234-123456789abc';

function server(options = {}) {
  const calls = { ai: [], search: [], writes: [], images: [] };
  let saved;
  const db = {
    async getSetting(key) {
      if (options.dbError) throw Error('private diagnostic');
      if (key.endsWith('conversation_chars')) return options.legacyLimit ?? 'broken';
      if (key === 'clinical_assistant.translate_provider') return 'libretranslate';
      return options.model || null;
    },
    async get(sql, params) {
      if (sql.includes('COUNT')) return { cnt: 0 };
      if (sql.includes('FROM clinical_assistant_chats') && params) return saved && params[0] === 1 && params[1] === 7 ? { id: params[0], title: saved[1], payload: saved[2] } : null;
      return saved && params[1] === 7 ? { id: 1, title: saved[1], payload: saved[2] } : null;
    },
    async run(sql, params) { calls.writes.push(sql); saved = params; return { lastInsertRowid: 1 }; },
    async query(sql, params) { return { rows: (params[0] || []).map(id => ({ id })) }; }
  };
  const ai = async (messages, settings) => {
    calls.ai.push({ messages, settings });
    return { content: 'Complete supported answer. [1]', finishReason: 'stop', model: 'synthetic' };
  };
  const source = { number: 1, title: 'Synthetic source', excerpt: 'Synthetic reference.', page: 9 };
  const mocks = {
    express, axios: { async post(url, payload) { if (String(url).includes('/translate')) return { data: { translatedText: 'Synthetic translation.' } }; calls.images.push(payload); return { data: { data: [{ b64_json: 'c3ludGhldGlj' }] } }; } }, crypto: require('node:crypto'), '../db/database': db,
    '../middleware/auth': { authMiddleware() {} },
    './auth': { __sendEmail: async () => false }, '../utils/ai': options.ai || { callAI: ai, callAIStream: ai },
    '../utils/errors': { gatewayUrl: path => 'http://synthetic.invalid' + path }, '../utils/litellm': { getLiteLLMHeaders: () => ({}) }, '../utils/logger': { audit() {}, error() {}, warn() {} },
    '../utils/crypto': { encryptString: value => 'encrypted:' + value, decryptString: value => value.replace(/^encrypted:/, '') },
    '../utils/redis': { async getJson() { return null; }, async setJson() {} }, '../utils/clinicalPromptPool': { createClinicalPromptPool: () => ({}) },
    '../utils/generatedImages': { ...require('../src/utils/generatedImages'), service: () => ({ async get() { return { jobId: 'synthetic', status: 'done', success: true }; }, async enqueue() { return { success: true, jobId: 'synthetic', status: 'pending' }; } }) },
    '../utils/generatedImageLinks': require('../src/utils/generatedImageLinks'),
    '../utils/imageTool': require('../src/utils/imageTool'),
    '../utils/clinicalPrompts': require('../src/utils/clinicalPrompts'),
    '../utils/clinicalMcpClient': {
      async semanticSearch(query) { calls.search.push(query); return {}; }, async getMcpHealth() { calls.health = calls.health || []; calls.health.push('health'); return {}; }
    },
    '../utils/clinicalRetrieval': {
      cleanSourceExcerpt: require('../src/utils/clinicalRetrieval').cleanSourceExcerpt,
      normalizeMcpSearchResponse: () => options.noSources ? [] : [source],
      normalizeMcpMultimodalResponse: () => [], dedupeSources: value => value,
      isVisualSourceQuery: () => false, classifyAndRerankMultimodalResults: async () => []
    },
    '../utils/clinicalAnswer': answer, '../utils/clinicalConversation': policy,
    '../utils/clinicalTranslation': require('../src/utils/clinicalTranslation'),
    '../utils/patientTakehome': require('../src/utils/patientTakehome')
  };
  const module = { exports: {} };
  vm.runInNewContext(read('src/routes/clinicalAssistant.js'), {
    module, exports: module.exports, console: quiet, Buffer, Map, TextEncoder,
    process: { env: { CLINICAL_ASSISTANT_MCP_WARMUP: 'false', CLINICAL_ASSISTANT_CONVERSATION_CHARS: options.limit, LITELLM_API_BASE: 'http://synthetic.invalid' } }, setTimeout() {},
    require(name) { assert.ok(Object.hasOwn(mocks, name), 'Unexpected import: ' + name); return mocks[name]; }
  });
  async function request(method, endpoint, body, userId = 7) {
    const handler = module.exports.stack.find(layer => layer.route && layer.route.path === endpoint && layer.route.methods[method]).route.stack.find(layer => layer.method === method).handle;
    const res = {
      statusCode: 200, headers: {}, events: '', body: null,
      status(code) { this.statusCode = code; return this; }, json(body) { this.body = body; return this; },
      setHeader(k, v) { this.headers[k] = v; }, write(text) { this.events += text; }, end() {}, flushHeaders() {}
    };
    await handler({ body, params: { id: 1 }, user: { id: userId } }, res);
    return res;
  }
  return { calls, request };
}

test('savedChatPayload persists bounded attachments and the generated image; invalid input is rejected before any write', () => {
  const messages = [{ role: 'user', content: 'Question', attachments: [png, jpeg] }, { role: 'assistant', content: 'Answer [1]', sources: [{ number: 1, title: 'S' }] }];
  const payload = policy.savedChatPayload({ messages, lastAnswer: 'Answer [1]', sources: [], generatedImage: asset });
  assert.deepEqual(payload.messages[0].attachments, [png, jpeg], 'normalized roundtrip incl. name');
  assert.equal(payload.generatedImage, asset, 'sidebar image now persists');
  assert.equal(payload.messages[1].attachments, undefined, 'absent attachments stay absent');

  const invalid = [
    [{ dataBase64: png.dataBase64, mimeType: 'image/svg+xml' }],
    [{ dataBase64: 'not base64!!!', mimeType: 'image/png' }],
    [{ dataBase64: canonical(5 * 1024 * 1024 + 1), mimeType: 'image/png' }],
    [png, jpeg, png, jpeg, png],
    'nope'
  ];
  for (const attachments of invalid) {
    assert.throws(() => policy.savedChatPayload({ messages: [{ role: 'user', content: 'Q', attachments }] }), error => error.statusCode === 400);
  }
  const longName = { dataBase64: png.dataBase64, mimeType: png.mimeType, name: 'n'.repeat(400) };
  const trimmed = policy.savedChatPayload({ messages: [{ role: 'user', content: 'Q', attachments: [longName] }] });
  assert.equal(trimmed.messages[0].attachments[0].name.length, 255, 'names clip at 255 characters');

  for (const generatedImage of ['data:image/svg+xml;base64,PHN2Zz4=', 'javascript:alert(1)']) {
    assert.throws(() => policy.savedChatPayload({ messages: [], generatedImage }), error => error.statusCode === 400);
  }
  assert.throws(() => policy.savedChatPayload({ messages: [{ role: 'user', content: 'x'.repeat(policy.MAX_SAVED_CHAT_BYTES) }] }), error => error.statusCode === 413);
});

test('POST /chats stores and reopens attachments; the same id updates in place instead of duplicating', async () => {
  const app = server();
  const body = { messages: [{ role: 'user', content: 'Question with image', attachments: [png] }, { role: 'assistant', content: 'Answer [1]' }], lastAnswer: 'Answer [1]', sources: [] };
  const saved = await app.request('post', '/clinical-assistant/chats', body);
  assert.equal(saved.statusCode, 200);
  const reopened = await app.request('get', '/clinical-assistant/chats/:id', {});
  assert.deepEqual(JSON.parse(JSON.stringify(reopened.body.chat.payload.messages[0].attachments)), [png]);
  const updated = await app.request('post', '/clinical-assistant/chats', { id: 1, messages: [{ role: 'user', content: 'Edited', attachments: [jpeg] }], lastAnswer: '' });
  assert.equal(updated.statusCode, 200);
  assert.equal(updated.body.id, 1);
  assert.ok(app.calls.writes.some(sql => /UPDATE clinical_assistant_chats/.test(sql)), 'update path used');
  const foreign = await app.request('post', '/clinical-assistant/chats', { id: 1, messages: [{ role: 'user', content: 'x' }] }, 8);
  assert.equal(foreign.statusCode, 404, 'updates are owner-bound');
  const missing = await app.request('post', '/clinical-assistant/chats', { id: 99, messages: [{ role: 'user', content: 'x' }] });
  assert.equal(missing.statusCode, 404);
  const badId = await app.request('post', '/clinical-assistant/chats', { id: 'not-a-number', messages: [{ role: 'user', content: 'x' }] });
  assert.equal(badId.statusCode, 400);
});

test('client restores attachments as thumbnails and saves them plus the generated image', t => {
  const dom = new JSDOM('<div id="assistant-tab">' + read('public/components/assistant.html') + '</div>', { url: 'https://example.test', runScripts: 'outside-only' });
  const window = dom.window;
  window.eval(read('public/js/accountBoundary.js'));
  assert.equal(window.AccountBoundary.enter({ id: 'synthetic-rendering-owner' }, true), true);
  window.marked = marked;
  window.DOMPurify = require('dompurify')(window);
  window.matchMedia = () => ({ matches: true });
  const saves = [];
  const context = { window, document: window.document, console, URL, Blob, TextDecoder, AbortController,
    setTimeout() {}, clearTimeout() {}, showToast() {}, EMPTY_PROMPT_SETS: [[]],
    createAssistantImageStore: () => ({ clear() {}, renderGeneratedImage: src => '<img src="' + src + '">' }),
    fetchSavedAssistantChats: async () => ({ success: true, chats: [] }),
    fetchAssistantImageJobs: async () => ({ success: true, jobs: [{ jobId: 'g1', imageUrl: asset }] }),
    saveAssistantChat: async body => { saves.push(JSON.parse(JSON.stringify(policy.savedChatPayload(body)))); return { success: true, id: 1 }; } };
  vm.createContext(context);
  for (const file of ['assistant/citations.js', 'assistant/sources.js', 'assistant/sharing.js', 'generatedImages.js', 'assistant/export.js', 'clinicalAssistant.js']) {
    vm.runInContext(read('public/js/' + file).replace(/^import[\s\S]*?from ['"][^'"]+['"];\s*/gm, '').replace(/^export /gm, ''), context);
  }
  t.after(() => window.close());
  context.bindEvents();
  context.appendMessage('user', 'Question with image', [], [], false, { attachments: [png] });
  let thumbs = window.document.querySelector('.assistant-msg.user .assistant-message-attachments');
  assert.ok(thumbs, 'thumbnail strip rendered');
  assert.ok(thumbs.querySelector('img').getAttribute('src').startsWith('data:image/png;base64,'));
  assert.match(thumbs.textContent, /chest\.png/);
  context.appendMessage('assistant', 'Answer [1].', [{ number: 1, title: 'S', page: 2 }]);
  context.lastAnswer = 'Answer [1].';
  context.lastGeneratedImageSrc = asset;
  context.scheduleAutosave();
  return context.performAutosave().then(() => {
    assert.deepEqual(saves[0].messages[0].attachments, [png]);
    assert.equal(saves[0].generatedImage, asset);
    context.restoreSavedChat(saves[0]);
    thumbs = window.document.querySelector('.assistant-msg.user .assistant-message-attachments');
    assert.ok(thumbs, 'attachments restored after reload');
    // The user's image library lives inside the Create image popup.
    window.document.getElementById('btn-assistant-create-image').click();
    var waitFor = function(remaining) {
      return new Promise(function(resolve) { setImmediate(resolve); }).then(function() {
        if (window.document.querySelector('#create-image-history img')) return Promise.resolve();
        if (remaining <= 0) return Promise.resolve();
        return waitFor(remaining - 1);
      });
    };
    return waitFor(12).then(function() {
      const galleryWrap = window.document.querySelector('#create-image-history');
      const galleryImg = window.document.querySelector('#create-image-history img');
      assert.ok(galleryImg && galleryImg.getAttribute('src') === asset, 'generated image restored into the image history popup; popup=' + (window.document.getElementById('assistant-create-image-modal') ? 'open' : 'missing') + ' history=' + (galleryWrap ? galleryWrap.innerHTML.slice(0, 120) : 'missing'));
    });
    assert.equal(context.messages[0].content, 'Question with image', 'raw transcript canonical');
  });
});