diff --git a/scripts/tests/test-qmd-daemon.mjs b/scripts/tests/test-qmd-daemon.mjs index accbb29..cbfee61 100644 --- a/scripts/tests/test-qmd-daemon.mjs +++ b/scripts/tests/test-qmd-daemon.mjs @@ -1,75 +1,170 @@ // scripts/tests/test-qmd-daemon.mjs import { test } from 'node:test'; import assert from 'node:assert/strict'; -import { createServer } from '../../shared/qmd/qmd-daemon.mjs'; +import { execFileSync } from 'node:child_process'; +import { mkdtempSync, rmSync, utimesSync, writeFileSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join, resolve } from 'node:path'; +import { createServer, gitFingerprint } from '../../shared/qmd/qmd-daemon.mjs'; import { collectionName } from '../../shared/qmd/qmd-lib.mjs'; -function fakeStore() { +const defaultHits = () => [{ displayPath: `${collectionName('/r')}/src/a.ts`, body: 'x\nretry here', bestChunkPos: 2, bestChunk: 'retry here', score: 0.91 }]; + +function fakeStore(updateGate, hits = defaultHits) { const calls = []; return { calls, - listCollections: async () => [{ name: 'notes', includeByDefault: true }], - getDefaultCollectionNames: async () => ['notes', 'repo-aaaaaaaaaaaa'], + listCollections: async () => [{ name: 'notes', includeByDefault: true }, ...calls.filter((c) => c[0] === 'add').map((c) => ({ name: c[1] }))], + getDefaultCollectionNames: async () => ['notes', 'repo-aaaaaaaaaaaa', 'repo-notes'], addCollection: async (name, opts) => calls.push(['add', name, opts]), - update: async (opts) => { calls.push(['update', opts]); return { indexed: 1 }; }, + update: async (opts) => { calls.push(['update', opts]); await updateGate; return { indexed: 1 }; }, embed: async (opts) => { calls.push(['embed', opts]); return { chunksEmbedded: 1 }; }, - search: async (opts) => { calls.push(['search', opts]); return [{ displayPath: `${collectionName('/r')}/src/a.ts`, body: 'x\nretry here', bestChunkPos: 2, bestChunk: 'retry here', score: 0.91, file: 'qmd://repo/src/a.ts' }]; }, + search: async (opts) => { calls.push(['search', opts]); return hits(); }, }; } -async function start(overrides = {}) { - const store = fakeStore(); +async function start(t, overrides = {}, { updateGate, hits } = {}) { + const store = fakeStore(updateGate, hits); let clock = 0; const idle = []; const daemon = createServer({ store, token: 't0k', idleMs: 1000, now: () => clock, onIdle: () => idle.push(clock), - extractSnippet: (body) => ({ line: 2, snippet: 'retry here' }), + extractSnippet: () => ({ line: 2, snippet: 'retry here' }), git: () => 'fingerprint-1', excludeCollection: async () => {}, recordCollection: () => {}, ...overrides, }); await new Promise((done) => daemon.server.listen(0, '127.0.0.1', done)); + t.after(() => { daemon.server.closeAllConnections(); daemon.server.close(); }); const port = daemon.server.address().port; const call = (path, body, token = 't0k') => fetch(`http://127.0.0.1:${port}${path}`, { - method: body === undefined ? 'GET' : 'POST', headers: { 'x-qmd-token': token, 'content-type': 'application/json' }, + method: body === undefined ? 'GET' : 'POST', headers: { connection: 'close', 'x-qmd-token': token, 'content-type': 'application/json' }, body: body === undefined ? undefined : JSON.stringify(body) }); - return { daemon, store, call, setClock: (v) => { clock = v; }, idle, close: () => daemon.server.close() }; + const raw = (path, body) => fetch(`http://127.0.0.1:${port}${path}`, { method: 'POST', headers: { connection: 'close', 'x-qmd-token': 't0k' }, body }); + return { daemon, store, call, raw, setClock: (v) => { clock = v; }, idle }; } -test('rejects requests without the token', async () => { - const s = await start(); +const wait = (ms) => new Promise((done) => setTimeout(done, ms)); +const gateOf = () => { let release; const gate = new Promise((done) => { release = done; }); return { gate, release }; }; + +test('rejects requests without the token', async (t) => { + const s = await start(t); assert.equal((await s.call('/health', undefined, 'wrong')).status, 403); - s.close(); }); -test('health identifies the service', async () => { - const s = await start(); +test('health identifies the service', async (t) => { + const s = await start(t); assert.equal((await (await s.call('/health')).json()).service, 'ai-config-qmd'); - s.close(); }); -test('search adds the repo collection once, refreshes on change, and filters other repos', async () => { - const s = await start(); +test('search adds the repo collection once, refreshes on change, and filters other repos', async (t) => { + const s = await start(t); const body = await (await s.call('/search', { repo: '/r', query: 'where are retries', limit: 3 })).json(); assert.deepEqual(body.results, [{ path: 'src/a.ts', line: 2, score: 0.91, snippet: 'retry here' }]); const search = s.store.calls.find((c) => c[0] === 'search')[1]; assert.equal(search.query, 'where are retries'); assert.ok(search.collections.includes('notes')); assert.ok(!search.collections.includes('repo-aaaaaaaaaaaa')); - assert.equal(search.collections.filter((n) => n.startsWith('repo-')).length, 1); + assert.ok(search.collections.includes('repo-notes'), 'only hash-named repo collections are filtered'); + assert.equal(search.collections.filter((n) => /^repo-[0-9a-f]{12}$/.test(n)).length, 1); assert.equal(s.store.calls.filter((c) => c[0] === 'add').length, 1); + assert.equal(s.store.calls.find((c) => c[0] === 'add')[2].path, resolve('/r')); await s.call('/search', { repo: '/r', query: 'again', limit: 3 }); assert.equal(s.store.calls.filter((c) => c[0] === 'update').length, 1, 'unchanged fingerprint skips update'); - s.close(); }); -test('rejects a missing query', async () => { - const s = await start(); +test('rejects a missing query', async (t) => { + const s = await start(t); assert.equal((await s.call('/search', { repo: '/r' })).status, 400); - s.close(); }); -test('idle check calls onIdle only after the idle window', async () => { - const s = await start(); +test('idle check calls onIdle only after the idle window', async (t) => { + const s = await start(t); s.setClock(500); s.daemon.checkIdle(); assert.deepEqual(s.idle, []); s.setClock(1500); s.daemon.checkIdle(); assert.deepEqual(s.idle, [1500]); - s.close(); +}); + +test('a changed fingerprint triggers update and embed again', async (t) => { + let fingerprint = 'one'; + const s = await start(t, { git: () => fingerprint }); + await s.call('/search', { repo: '/r', query: 'q' }); + fingerprint = 'two'; + await s.call('/search', { repo: '/r', query: 'q' }); + assert.equal(s.store.calls.filter((c) => c[0] === 'update').length, 2); + assert.equal(s.store.calls.filter((c) => c[0] === 'embed').length, 2); +}); + +test('fingerprint changes when an already dirty file is edited again', (t) => { + const dir = mkdtempSync(join(tmpdir(), 'qmd-fp-')); + t.after(() => rmSync(dir, { recursive: true, force: true })); + const git = (...args) => execFileSync('git', ['-C', dir, '-c', 'user.name=t', '-c', 'user.email=t@t', ...args], { stdio: 'ignore' }); + git('init', '-q'); + writeFileSync(join(dir, 'a.ts'), 'one'); + git('add', '.'); + git('commit', '-qm', 'init'); + writeFileSync(join(dir, 'a.ts'), 'two'); + const first = gitFingerprint(dir); + writeFileSync(join(dir, 'a.ts'), 'three!'); + utimesSync(join(dir, 'a.ts'), new Date(), new Date(Date.now() + 5000)); + assert.notEqual(gitFingerprint(dir), first); +}); + +test('a failing exclude does not stop the search and is retried', async (t) => { + let attempts = 0; + const recorded = []; + const s = await start(t, { + excludeCollection: async () => { attempts++; if (attempts === 1) throw new Error('boom'); }, + recordCollection: (name) => recorded.push(name), + }); + const first = await s.call('/search', { repo: '/r', query: 'q' }); + assert.equal(first.status, 200); + assert.equal((await first.json()).results.length, 1); + assert.deepEqual(recorded, []); + await s.call('/search', { repo: '/r', query: 'q' }); + assert.equal(attempts, 2); + assert.equal(recorded.length, 1); +}); + +test('concurrent searches share one refresh', async (t) => { + const { gate, release } = gateOf(); + const s = await start(t, {}, { updateGate: gate }); + const both = Promise.all([s.call('/search', { repo: '/r', query: 'a' }), s.call('/search', { repo: '/r', query: 'b' })]); + await wait(100); + release(); + await both; + assert.equal(s.store.calls.filter((c) => c[0] === 'update').length, 1); +}); + +test('idle check waits for in-flight requests', async (t) => { + const { gate, release } = gateOf(); + const s = await start(t, {}, { updateGate: gate }); + const pending = s.call('/search', { repo: '/r', query: 'a' }); + await wait(100); + s.setClock(5000); s.daemon.checkIdle(); assert.deepEqual(s.idle, []); + release(); + await pending; +}); + +test('malformed or non-object bodies give 400 and limit is clamped', async (t) => { + const s = await start(t); + assert.equal((await s.raw('/search', '{nope')).status, 400); + assert.equal((await s.raw('/search', '[1]')).status, 400); + assert.equal((await s.raw('/search', 'null')).status, 400); + await s.call('/search', { repo: '/r', query: 'q', limit: 9999 }); + await s.call('/search', { repo: '/r', query: 'q', limit: 'x' }); + const limits = s.store.calls.filter((c) => c[0] === 'search').map((c) => c[1].limit); + assert.deepEqual(limits, [20, 5]); +}); + +test('non-repo hits become qmd:// URIs', async (t) => { + const s = await start(t, {}, { hits: () => [{ displayPath: 'notes/todo.md', body: 'x', bestChunkPos: 0, bestChunk: 'x', score: 0.5 }] }); + const body = await (await s.call('/search', { repo: '/r', query: 'q' })).json(); + assert.equal(body.results[0].path, 'qmd://notes/todo.md'); +}); + +test('refresh, stop, 404 and 403 on POST routes', async (t) => { + const s = await start(t); + assert.deepEqual(await (await s.call('/refresh', { repo: '/r' })).json(), { collection: collectionName('/r'), updated: true }); + assert.equal((await s.call('/nothing', { repo: '/r' })).status, 404); + assert.equal((await s.call('/search', { repo: '/r', query: 'q' }, 'wrong')).status, 403); + assert.deepEqual(await (await s.call('/stop', {})).json(), { stopping: true }); + assert.equal(s.idle.length, 1); }); diff --git a/shared/qmd/qmd-daemon.mjs b/shared/qmd/qmd-daemon.mjs index 04958e1..78ee4d9 100644 --- a/shared/qmd/qmd-daemon.mjs +++ b/shared/qmd/qmd-daemon.mjs @@ -1,38 +1,66 @@ // shared/qmd/qmd-daemon.mjs — keeps QMD's models loaded for fast searches; exits when idle. import { execFileSync } from 'node:child_process'; -import { randomBytes } from 'node:crypto'; +import { randomBytes, timingSafeEqual } from 'node:crypto'; +import { statSync } from 'node:fs'; import { createServer as createHttpServer } from 'node:http'; -import { join } from 'node:path'; +import { join, resolve } from 'node:path'; import { pathToFileURL } from 'node:url'; import { MASK, collectionName, idleMinutes, importQmd, qmdConfigPath, qmdDbPath, readJson, runQmdCli, stateDir, writePrivateJson } from './qmd-lib.mjs'; const MAX_BODY = 64 * 1024; -function gitFingerprint(repo) { +// Dirty files keep the same porcelain line when edited again, so mtime and size of each listed path are part of the fingerprint. +export function gitFingerprint(repo) { try { - return execFileSync('git', ['-C', repo, 'status', '--porcelain=v1', '-uall'], { encoding: 'utf8', maxBuffer: 16 * 1024 * 1024, windowsHide: true }) - + execFileSync('git', ['-C', repo, 'rev-parse', 'HEAD'], { encoding: 'utf8', windowsHide: true }); + const status = execFileSync('git', ['-C', repo, 'status', '--porcelain=v1', '-uall', '-z'], { encoding: 'utf8', maxBuffer: 16 * 1024 * 1024, windowsHide: true }); + const stats = status.split('\0').filter((entry) => entry.length > 3).map((entry) => { + try { + const stat = statSync(join(repo, entry.slice(3))); + return `${entry}|${stat.mtimeMs}:${stat.size}`; + } catch { + return `${entry}|deleted`; + } + }); + return `${stats.join('\n')}\n` + execFileSync('git', ['-C', repo, 'rev-parse', 'HEAD'], { encoding: 'utf8', windowsHide: true }); } catch { return String(Date.now()); } } +const REPO_COLLECTION = /^repo-[0-9a-f]{12}$/; + +function tokenMatches(given, expected) { + const a = Buffer.from(String(given ?? '')); + const b = Buffer.from(expected); + return a.length === b.length && timingSafeEqual(a, b); +} + export function createServer({ store, token, idleMs, now = Date.now, onIdle, extractSnippet, git = gitFingerprint, excludeCollection, recordCollection }) { let lastRequest = now(); + let active = 0; const known = new Set(); const fingerprints = new Map(); - async function refresh(repo) { + const inFlight = new Map(); + + function refresh(repo) { const name = collectionName(repo); + if (!inFlight.has(name)) inFlight.set(name, doRefresh(repo, name).finally(() => inFlight.delete(name))); + return inFlight.get(name); + } + + async function doRefresh(repo, name) { if (!known.has(name)) { const existing = (await store.listCollections()).some((c) => c.name === name); - if (!existing) { - await store.addCollection(name, { path: repo, pattern: MASK }); + if (!existing) await store.addCollection(name, { path: resolve(repo), pattern: MASK }); + try { await excludeCollection(name); recordCollection(name); + known.add(name); + } catch (error) { + console.error(`qmd-daemon: could not exclude ${name} from default search: ${error.message}`); } - known.add(name); } const fingerprint = git(repo); if (fingerprints.get(name) === fingerprint) return { collection: name, updated: false }; @@ -44,8 +72,8 @@ export function createServer({ store, token, idleMs, now = Date.now, onIdle, ext async function search({ repo, query, limit }) { const { collection } = await refresh(repo); - const defaults = (await store.getDefaultCollectionNames()).filter((n) => !n.startsWith('repo-')); - const results = await store.search({ query, collections: [collection, ...defaults], limit: limit || 5, chunkStrategy: 'auto' }); + const defaults = (await store.getDefaultCollectionNames()).filter((n) => !REPO_COLLECTION.test(n)); + const results = await store.search({ query, collections: [collection, ...defaults], limit: Number.isInteger(limit) ? Math.min(20, Math.max(1, limit)) : 5, chunkStrategy: 'auto' }); return results.map((r) => { const { line, snippet } = extractSnippet(r.body, query, 300, r.bestChunkPos, r.bestChunk.length); // displayPath is /: repo hits become repo-relative, others stay addressable as qmd:// URIs. @@ -59,16 +87,20 @@ export function createServer({ store, token, idleMs, now = Date.now, onIdle, ext response.writeHead(status, { 'content-type': 'application/json' }); response.end(JSON.stringify(value)); }; - if (request.headers['x-qmd-token'] !== token) return reply(403, { error: 'forbidden' }); + if (!tokenMatches(request.headers['x-qmd-token'], token)) return reply(403, { error: 'forbidden' }); + active++; lastRequest = now(); try { if (request.method === 'GET' && request.url === '/health') return reply(200, { status: 'ok', service: 'ai-config-qmd', pid: process.pid }); + request.setEncoding('utf8'); let raw = ''; for await (const chunk of request) { raw += chunk; if (raw.length > MAX_BODY) return reply(413, { error: 'body too large' }); } - const body = raw ? JSON.parse(raw) : {}; + let body; + try { body = raw ? JSON.parse(raw) : {}; } catch { return reply(400, { error: 'invalid JSON' }); } + if (!body || typeof body !== 'object' || Array.isArray(body)) return reply(400, { error: 'body must be an object' }); if (request.method === 'POST' && request.url === '/stop') { reply(200, { stopping: true }); return onIdle(); @@ -83,11 +115,12 @@ export function createServer({ store, token, idleMs, now = Date.now, onIdle, ext } catch (error) { return reply(500, { error: String(error?.message || error) }); } finally { + active--; lastRequest = now(); } }); - const checkIdle = () => { if (now() - lastRequest >= idleMs) onIdle(); }; + const checkIdle = () => { if (active === 0 && now() - lastRequest >= idleMs) onIdle(); }; return { server, refresh, search, checkIdle }; } @@ -111,14 +144,14 @@ async function main() { stopping = true; const latest = readJson(daemonFile, null); if (latest?.pid === process.pid) writePrivateJson(daemonFile, { stopped: true }); + daemon.server.close(); + daemon.server.closeIdleConnections(); await store.close(); process.exit(0); }; const daemon = createServer({ store, token, idleMs: idleMinutes() * 60_000, onIdle: stop, extractSnippet, - excludeCollection: async (name) => { - try { runQmdCli(['collection', 'exclude', name], { windowsHide: true, stdio: 'ignore' }); } catch { /* indexing still works without the exclusion */ } - }, + excludeCollection: async (name) => { runQmdCli(['collection', 'exclude', name], { windowsHide: true, stdio: 'ignore' }); }, recordCollection: (name) => { const names = new Set(readJson(collectionsFile, [])); names.add(name);