pediatric-ai-scribe-v3/test/clinical-mcp-session-lifecycle.test.js

const { test } = require('node:test');
const assert = require('node:assert/strict');
const Module = require('node:module');
const { Readable } = require('node:stream');

function deferred() {
  let resolve;
  let reject;
  const promise = new Promise((yes, no) => { resolve = yes; reject = no; });
  return { promise, resolve, reject };
}

function loadClient(overrides = {}) {
  const calls = [];
  let count = 0;
  const axios = {
    post: async (url, payload, config) => {
      const call = { method: 'POST', url, payload, config };
      calls.push(call);
      const init = payload.method === 'initialize';
      const response = {
        status: init ? 200 : 202, config: { url },
        headers: init ? { 'mcp-session-id': 'session-' + (++count) } : {},
        data: payload.method === 'notifications/initialized' ? '' : JSON.stringify({ jsonrpc: '2.0', result: { content: [] } })
      };
      return overrides.post ? overrides.post(call, response) : response;
    },
    delete: async (url, config) => {
      const call = { method: 'DELETE', url, config };
      calls.push(call);
      return overrides.delete ? overrides.delete(call) : { status: 204 };
    }
  };
  const target = require.resolve('../src/utils/clinicalMcpClient');
  delete require.cache[target];
  const original = Module._load;
  Module._load = function(request) {
    if (request === 'axios') return axios;
    return original.apply(this, arguments);
  };
  let client;
  try { client = require(target); } finally { Module._load = original; }
  return { client, calls, deletes: () => calls.filter(c => c.method === 'DELETE').map(c => c.config.headers['mcp-session-id']) };
}

const method = c => c.payload && c.payload.method;
const turn = () => new Promise(resolve => setImmediate(resolve));

test('live session reuse, real initialized notification and serialized tool calls', async () => {
  const held = deferred();
  const started = deferred();
  let tools = 0;
  const { client, calls, deletes } = loadClient({ post: async (call, response) => {
    if (method(call) === 'initialize') response.data = Readable.from([response.data]);
    if (method(call) === 'tools/call' && ++tools === 1) { started.resolve(); await held.promise; }
    return response;
  } });
  const first = client.semanticSearch('one', { limit: 4 });
  await started.promise;
  const second = client.indexedTopicSuggestions(8); // any second tool call proves serialization
  await turn();
  assert.equal(tools, 1);
  held.resolve();
  await Promise.all([first, second]);
  await client.semanticSearch('three', { limit: 4 });
  assert.deepEqual(calls.filter(c => c.method === 'POST').map(method), ['initialize', 'notifications/initialized', 'tools/call', 'tools/call', 'tools/call']);
  const notice = calls[1];
  assert.equal(notice.payload.jsonrpc, '2.0');
  assert.equal(Object.hasOwn(notice.payload, 'id'), false);
  assert.equal(notice.config.headers['mcp-session-id'], 'session-1');
  assert.equal(notice.url, calls[0].url);
  assert.deepEqual(deletes(), []);
  await client.closeMcpSession();
  assert.deepEqual(deletes(), ['session-1']);
});

test('expired session is discarded even if replacement initialize fails', async t => {
  let now = 0;
  t.mock.method(Date, 'now', () => now);
  let initializations = 0;
  const { client, deletes } = loadClient({ post: async (call, response) => {
    if (method(call) === 'initialize' && ++initializations === 2) throw new Error('failed replacement');
    return response;
  } });
  await client.warmMcpSession();
  now = 600001;
  await assert.rejects(client.semanticSearch('two'));
  assert.deepEqual(deletes(), ['session-1']);
  await client.closeMcpSession();
});

for (const failure of ['malformed', 'rpc', 'empty', 'body', 'notification', 'notification-rpc', 'http']) {
  test('known initialize session is unwound on ' + failure + ' failure', async () => {
    const { client, calls, deletes } = loadClient({ post: async (call, response) => {
      if (method(call) === 'initialize') {
        if (failure === 'malformed') response.data = '{broken';
        if (failure === 'rpc') response.data = '{"error":{"message":"secret-session-1"}}';
        if (failure === 'empty') response.data = '';
        if (failure === 'body') response.data = Readable.from((async function*() { throw new Error('secret-body'); })());
        if (failure === 'http') throw Object.assign(new Error('secret-http'), { response: { ...response, status: 503 } });
      }
      if (method(call) === 'notifications/initialized') {
        if (failure === 'notification') throw new Error('secret-notification');
        if (failure === 'notification-rpc') response.data = '{"error":{"message":"secret-notification"}}';
      }
      return response;
    } });
    await assert.rejects(client.warmMcpSession(), e => !/secret/.test(e.message));
    assert.deepEqual(deletes(), ['session-1']);
    assert.equal(calls.filter(c => method(c) === 'initialize').length, 1);
    await client.closeMcpSession();
    assert.deepEqual(deletes(), ['session-1']);
  });
}

test('initialized notification never falls back to another session endpoint', async t => {
  const previous = process.env.CLINICAL_ASSISTANT_MCP_INITIALIZE_TIMEOUT_MS;
  process.env.CLINICAL_ASSISTANT_MCP_INITIALIZE_TIMEOUT_MS = '40';
  t.after(() => {
    if (previous === undefined) delete process.env.CLINICAL_ASSISTANT_MCP_INITIALIZE_TIMEOUT_MS;
    else process.env.CLINICAL_ASSISTANT_MCP_INITIALIZE_TIMEOUT_MS = previous;
  });
  t.mock.method(console, 'warn', () => {});
  let owner;
  const { client, calls, deletes } = loadClient({ post: async (call, response) => {
    if (method(call) === 'initialize') owner = call.url;
    if (method(call) === 'notifications/initialized' && call.url === owner) {
      throw Object.assign(new Error('temporarily unavailable'), { response: { status: 503 } });
    }
    return response; // An alternate endpoint would accept, but must never be contacted.
  } });
  await assert.rejects(client.warmMcpSession());
  const notices = calls.filter(c => method(c) === 'notifications/initialized');
  assert.ok(notices.length > 0);
  assert.ok(notices.every(c => c.url === owner));
  assert.deepEqual(deletes(), ['session-1']);
  await client.closeMcpSession();
});

test('ambiguous initialize transport failure is not retried', async () => {
  const { client, calls, deletes } = loadClient({ post: async () => {
    throw Object.assign(new Error('secret-transport'), { code: 'ECONNRESET' });
  } });
  await assert.rejects(client.warmMcpSession(), /MCP request failed/);
  assert.equal(calls.length, 1);
  assert.deepEqual(deletes(), []);
  await client.closeMcpSession();
});

test('invalid-session recovery discards the failed object, not a newer warm session', async t => {
  let now = 0;
  t.mock.method(Date, 'now', () => now);
  const held = deferred();
  const started = deferred();
  let tools = 0;
  const { client, calls, deletes } = loadClient({ post: async (call, response) => {
    if (method(call) === 'tools/call' && ++tools === 1) {
      started.resolve(); await held.promise;
      throw Object.assign(new Error('invalid'), { response: { status: 404 } });
    }
    return response;
  } });
  const search = client.semanticSearch('one');
  await started.promise;
  now = 600001;
  await client.warmMcpSession();
  held.resolve();
  await search;
  assert.deepEqual(deletes(), ['session-1']);
  assert.equal(calls.filter(c => method(c) === 'initialize').length, 2);
  assert.equal(calls.filter(c => method(c) === 'tools/call')[1].config.headers['mcp-session-id'], 'session-2');
  await client.closeMcpSession();
});

test('both invalid tool attempts are discarded, ordinary tool errors keep the shared session', async () => {
  let invalid = true;
  const { client, deletes } = loadClient({ post: async (call, response) => {
    if (method(call) === 'tools/call') throw Object.assign(new Error('tool failure'), invalid ? { response: { status: 404 } } : {});
    return response;
  } });
  await assert.rejects(client.semanticSearch('one'));
  assert.deepEqual(deletes(), ['session-1', 'session-2']);
  invalid = false;
  await assert.rejects(client.semanticSearch('two'));
  assert.deepEqual(deletes(), ['session-1', 'session-2']);
  await client.closeMcpSession();
  assert.deepEqual(deletes(), ['session-1', 'session-2', 'session-3']);
});

test('duplicate shutdown rejects warm/new/queued calls and unwinds late initialize', async () => {
  const held = deferred();
  const started = deferred();
  const { client, calls, deletes } = loadClient({ post: async (call, response) => {
    if (method(call) === 'initialize') { started.resolve(); await held.promise; }
    return response;
  } });
  const warm = client.warmMcpSession();
  const rejectedWarm = assert.rejects(warm);
  await started.promise;
  const queued = assert.rejects(client.semanticSearch('queued'));
  const closing = client.closeMcpSession();
  assert.equal(client.closeMcpSession(), closing);
  await assert.rejects(client.warmMcpSession());
  await assert.rejects(client.semanticSearch('new'));
  held.resolve();
  await Promise.all([closing, rejectedWarm, queued]);
  assert.deepEqual(deletes(), ['session-1']);
  assert.equal(calls.filter(c => method(c) === 'tools/call').length, 0);
});

test('shutdown during initialized notification cannot cache a late session', async () => {
  const held = deferred();
  const started = deferred();
  const { client, deletes } = loadClient({ post: async (call, response) => {
    if (method(call) === 'notifications/initialized') { started.resolve(); await held.promise; }
    return response;
  } });
  const warming = assert.rejects(client.warmMcpSession());
  await started.promise;
  const closing = client.closeMcpSession();
  held.resolve();
  await Promise.all([warming, closing]);
  assert.deepEqual(deletes(), ['session-1']);
  await assert.rejects(client.warmMcpSession());
});

test('cleanup rejection does not lose an in-flight successful search', async () => {
  const held = deferred();
  const started = deferred();
  const { client, deletes } = loadClient({ post: async (call, response) => {
    if (method(call) === 'tools/call') { started.resolve(); await held.promise; }
    return response;
  }, delete: async () => { throw new Error('secret-cleanup'); } });
  const search = client.semanticSearch('one');
  await started.promise;
  const queued = assert.rejects(client.semanticSearch('queued'));
  const closing = client.closeMcpSession();
  held.resolve();
  assert.deepEqual(await search, { content: [] });
  await Promise.all([queued, closing]);
  assert.deepEqual(deletes(), ['session-1']);
});

test('DELETE and shutdown are bounded even when transport ignores abort or holds a body', async t => {
  t.mock.timers.enable({ apis: ['setTimeout'] });
  let config;
  const { client } = loadClient({ delete: call => { config = call.config; return new Promise(() => {}); } });
  await client.warmMcpSession();
  const closing = client.closeMcpSession();
  assert.equal(config.timeout, 5000);
  assert.equal(config.signal.aborted, false);
  t.mock.timers.tick(5000);
  await closing;
  assert.equal(config.signal.aborted, true);
});

test('shutdown bound also covers unknown in-flight initialization', async t => {
  t.mock.timers.enable({ apis: ['setTimeout'] });
  const held = deferred();
  const started = deferred();
  const { client, deletes } = loadClient({ post: async (call, response) => {
    started.resolve(); await held.promise; return response;
  } });
  const warm = assert.rejects(client.warmMcpSession());
  await started.promise;
  const closing = client.closeMcpSession();
  t.mock.timers.tick(5000);
  await closing;
  held.resolve();
  await warm;
  assert.deepEqual(deletes(), ['session-1']);
});

test('expired replacement succeeds even when old-session DELETE fails', async t => {
  let now = 0;
  t.mock.method(Date, 'now', () => now);
  const { client, calls, deletes } = loadClient({ delete: async () => { throw new Error('cleanup failed'); } });
  await client.semanticSearch('one');
  now = 600001;
  assert.deepEqual(await client.semanticSearch('two'), { content: [] });
  assert.deepEqual(deletes(), ['session-1']);
  assert.equal(calls.filter(c => method(c) === 'initialize').length, 2);
  await client.closeMcpSession();
});

test('connection refusal preserves URL fallback and the chosen URL owns handshake/tools/DELETE', async t => {
  const warnings = [];
  t.mock.method(console, 'warn', (...args) => warnings.push(args.join(' ')));
  let first = true;
  const { client, calls } = loadClient({ post: async (call, response) => {
    if (first) { first = false; throw Object.assign(new Error('secret endpoint'), { code: 'ECONNREFUSED' }); }
    return response;
  } });
  await client.semanticSearch('one');
  await client.closeMcpSession();
  assert.notEqual(calls[0].url, calls[1].url);
  assert.ok(calls.slice(1).every(call => call.url === calls[1].url));
  assert.deepEqual(warnings, ['[clinical-assistant] MCP endpoint unavailable']);
});

for (const scenario of ['broken initialize body', 'held DELETE body']) {
  test('real Axios against synthetic loopback: ' + scenario, { timeout: 10000 }, async t => {
    const http = require('node:http');
    const requests = [];
    const deleteClosed = deferred();
    const server = http.createServer(async (req, res) => {
      const chunks = [];
      for await (const chunk of req) chunks.push(chunk);
      const payload = chunks.length && JSON.parse(Buffer.concat(chunks).toString());
      requests.push({ method: req.method, payload, id: req.headers['mcp-session-id'] });
      if (req.method === 'DELETE') {
        res.on('close', () => deleteClosed.resolve());
        if (scenario === 'held DELETE body') { res.writeHead(503); res.write('held'); }
        else { res.writeHead(204); res.end(); }
      } else if (payload.method === 'initialize') {
        res.setHeader('mcp-session-id', 'synthetic-session');
        if (scenario === 'broken initialize body') {
          res.write('{');
          // Simulate a connection breaking after headers have reached the client.
          setTimeout(() => res.destroy(), 20);
        } else res.end('{"jsonrpc":"2.0","result":{"protocolVersion":"2024-11-05"}}');
      } else if (payload.method === 'notifications/initialized') { res.writeHead(202); res.end(); }
      else res.end('{"result":{"content":[]}}');
    });
    await new Promise(resolve => server.listen(0, '127.0.0.1', resolve));
    const previous = process.env.CLINICAL_ASSISTANT_MCP_URLS;
    process.env.CLINICAL_ASSISTANT_MCP_URLS = 'http://127.0.0.1:' + server.address().port + '/mcp';
    t.after(async () => {
      if (previous === undefined) delete process.env.CLINICAL_ASSISTANT_MCP_URLS;
      else process.env.CLINICAL_ASSISTANT_MCP_URLS = previous;
      server.closeAllConnections();
      await new Promise(resolve => server.close(resolve));
    });
    const target = require.resolve('../src/utils/clinicalMcpClient');
    delete require.cache[target];
    const client = require(target);
    if (scenario === 'broken initialize body') {
      await assert.rejects(client.warmMcpSession());
      assert.deepEqual(requests.map(r => r.method), ['POST', 'DELETE']);
    } else {
      assert.deepEqual(await client.semanticSearch('synthetic'), { content: [] });
      assert.deepEqual(requests.map(r => r.payload.method), ['initialize', 'notifications/initialized', 'tools/call']);
      const started = Date.now();
      await client.closeMcpSession();
      assert.ok(Date.now() - started < 6500, 'cleanup must not wait for the held body');
    }
    await deleteClosed.promise;
    assert.equal(requests.at(-1).id, 'synthetic-session');
    await client.closeMcpSession();
  });
}

test('real Axios tool deadline aborts a continuously trickling response', { timeout: 3000 }, async t => {
  const http = require('node:http');
  const bodyClosed = deferred();
  let completed = false;
  const server = http.createServer(async (req, res) => {
    const chunks = [];
    for await (const chunk of req) chunks.push(chunk);
    const payload = chunks.length && JSON.parse(Buffer.concat(chunks).toString());
    if (req.method === 'DELETE') { res.writeHead(204); res.end(); }
    else if (payload.method === 'initialize') {
      res.setHeader('mcp-session-id', 'trickle-session');
      res.end('{"result":{"protocolVersion":"2024-11-05"}}');
    } else if (payload.method === 'notifications/initialized') { res.writeHead(202); res.end(); }
    else {
      res.write('{"result":{"content":[');
      const interval = setInterval(() => res.write(' '), 15);
      const finish = setTimeout(() => { completed = true; res.end(']}}'); }, 450);
      res.on('close', () => { clearInterval(interval); clearTimeout(finish); bodyClosed.resolve(); });
    }
  });
  await new Promise(resolve => server.listen(0, '127.0.0.1', resolve));
  const keys = ['CLINICAL_ASSISTANT_MCP_URLS', 'CLINICAL_ASSISTANT_MCP_REQUEST_TIMEOUT_MS'];
  const previous = keys.map(key => process.env[key]);
  process.env[keys[0]] = 'http://127.0.0.1:' + server.address().port + '/mcp';
  process.env[keys[1]] = '100';
  let client;
  t.after(async () => {
    keys.forEach((key, i) => {
      if (previous[i] === undefined) delete process.env[key];
      else process.env[key] = previous[i];
    });
    if (client) await client.closeMcpSession();
    server.closeAllConnections();
    await new Promise(resolve => server.close(resolve));
  });
  const target = require.resolve('../src/utils/clinicalMcpClient');
  delete require.cache[target];
  client = require(target);
  await assert.rejects(client.semanticSearch('synthetic'), /MCP request failed/);
  await bodyClosed.promise;
  assert.equal(completed, false, 'deadline closes the socket before the trickling body completes');
});

test('shutdown between cached-session lookup and tool dispatch prevents a late call', async () => {
  const { client, calls } = loadClient();
  await client.warmMcpSession();
  const search = assert.rejects(client.semanticSearch('late'));
  await Promise.resolve(); // queued call entered getMcpSession, dispatch has not resumed
  await client.closeMcpSession();
  await search;
  assert.equal(calls.filter(c => method(c) === 'tools/call').length, 0);
});