diff --git a/scripts/lib/qmd.ps1 b/scripts/lib/qmd.ps1 index 198c715..379feff 100644 --- a/scripts/lib/qmd.ps1 +++ b/scripts/lib/qmd.ps1 @@ -1,4 +1,3 @@ -# scripts/lib/qmd.ps1 function Sync-Qmd { param([Parameter(Mandatory=$true)][string]$RepositoryRoot,[Parameter(Mandatory=$true)][string]$HomePath,[ValidateSet('Codex','Claude','Both')][Parameter(Mandatory=$true)][string]$Client,[Parameter(Mandatory=$true)][bool]$Enabled,[switch]$DryRun,[switch]$Update,[switch]$Summary) $arguments = @((Join-Path $PSScriptRoot 'qmd.py'), 'sync', '--root', $RepositoryRoot, '--home', $HomePath, '--client', $Client.ToLowerInvariant(), '--enabled', $Enabled.ToString().ToLowerInvariant()) diff --git a/scripts/lib/qmd.py b/scripts/lib/qmd.py index c8a2999..f2d6350 100644 --- a/scripts/lib/qmd.py +++ b/scripts/lib/qmd.py @@ -17,10 +17,10 @@ DEVICE_TIMEOUT = 600 WINDOWS = os.name == 'nt' -def run(arguments, env=None, timeout=1800): +def run(arguments, env=None, timeout=1800, check=True): result = subprocess.run(arguments, env=env, capture_output=True, text=True, encoding='utf-8', errors='replace', timeout=timeout, shell=WINDOWS and arguments[0] in ('npm', 'qmd')) - if result.returncode != 0: + if check and result.returncode != 0: raise ValueError(f'{" ".join(arguments[:3])} failed: {(result.stderr or result.stdout).strip()[-300:]}') return result.stdout @@ -63,7 +63,7 @@ class Setup: def load(self): if not self.state_path.exists(): - return {'version': 1, 'clients': [], 'device': None, 'owned_package': False, 'owned_models': [], 'models_block': False, 'files': {}} + return {'version': 1, 'clients': [], 'device': None, 'owned_package': False, 'owned_models': [], 'models_block': False, 'previous_models': None, 'files': {}} state = json.loads(self.state_path.read_text(encoding='utf-8')) if state.get('version') != 1 or state.get('device') not in (None, 'gpu', 'cpu'): raise ValueError('Invalid QMD setup state.') @@ -73,7 +73,12 @@ class Setup: return self.run(['node', str(self.scripts / script), *arguments], env=env, timeout=timeout) def package_installed(self): - listing = json.loads(self.run(['npm', 'list', '--global', '--depth=0', '--json']) or '{}') + # npm list exits non-zero on tree problems (ELSPROBLEMS) but still prints valid JSON. + output = self.run(['npm', 'list', '--global', '--depth=0', '--json'], check=False) or '{}' + try: + listing = json.loads(output) + except json.JSONDecodeError: + raise ValueError('npm list did not return JSON.') return PACKAGE in (listing.get('dependencies') or {}) def require_node(self): @@ -87,7 +92,8 @@ class Setup: if legacy_state.exists() or legacy.exists(): if legacy.is_symlink(): raise ValueError('Legacy semantic retrieval directory must not be a symbolic link.') - shutil.rmtree(legacy, ignore_errors=True) + if legacy.exists(): + shutil.rmtree(legacy) legacy_state.unlink(missing_ok=True) self.say('Removed the previous Python semantic retrieval installation.') ledger_path = self.base / 'extensions.json' @@ -95,14 +101,19 @@ class Setup: ledger = json.loads(ledger_path.read_text(encoding='utf-8')) owned = [r for r in ledger.get('resources', []) if r.get('kind') == 'qmd'] if owned: + state['owned_package'] = True + atomic_json(self.state_path, state) # record ownership before the ledger forgets it ledger['resources'] = [r for r in ledger['resources'] if r.get('kind') != 'qmd'] atomic_json(ledger_path, ledger) - state['owned_package'] = True def install_scripts(self, state): + if self.scripts.is_symlink(): + raise ValueError('QMD script directory must not be a symbolic link.') self.scripts.mkdir(parents=True, exist_ok=True) for name in SCRIPTS: target = self.scripts / name + if target.is_symlink(): + raise ValueError(f'QMD script must not be a symbolic link: {name}') expected = state['files'].get(name) if target.exists() and expected and digest(target) != expected: raise ValueError(f'Setup-owned QMD script changed: {name}') @@ -119,7 +130,12 @@ class Setup: self.install_scripts(state) atomic_json(self.state_path, state) if not state['models_block']: - self.node('qmd-config.mjs', 'set-models') + output = self.node('qmd-config.mjs', 'set-models') + try: + previous = json.loads(output.strip().splitlines()[-1])['previous'] + except (IndexError, KeyError, TypeError, json.JSONDecodeError): + raise ValueError('qmd-config set-models returned no previous models.') + state['previous_models'] = previous if isinstance(previous, dict) else None state['models_block'] = True atomic_json(self.state_path, state) before = {p.name for p in self.models.iterdir()} if self.models.is_dir() else set() @@ -159,10 +175,14 @@ class Setup: if re.fullmatch(r'repo-[0-9a-f]{12}', name): try: self.run(['qmd', 'collection', 'remove', name]) - except ValueError: + except (ValueError, subprocess.SubprocessError): pass if state['models_block'] and self.scripts.joinpath('qmd-config.mjs').exists(): - self.node('qmd-config.mjs', 'unset-models') + previous = state.get('previous_models') + if previous is not None: + self.node('qmd-config.mjs', 'restore-models', json.dumps(previous)) + else: + self.node('qmd-config.mjs', 'unset-models') for name in state['owned_models']: path = self.models / name if name == Path(name).name and path.is_file() and not path.is_symlink(): @@ -178,6 +198,9 @@ class Setup: if not self.state_path.exists(): if not self.dry_run: self.migrate(state) + if state['owned_package']: + self.run(['npm', 'uninstall', '--global', PACKAGE]) + self.state_path.unlink(missing_ok=True) return if self.dry_run: self.say('DRYRUN disable QMD local search') @@ -203,7 +226,9 @@ class Setup: env['QMD_LLAMA_GPU' if device == 'gpu' else 'QMD_FORCE_CPU'] = 'auto' if device == 'gpu' else '1' try: result = json.loads(self.node('qmd-benchmark.mjs', '--device', device, env=env, timeout=DEVICE_TIMEOUT)) - except (ValueError, subprocess.TimeoutExpired) as error: + if not isinstance(result, dict): + raise ValueError('benchmark returned a non-object result') + except (ValueError, subprocess.SubprocessError) as error: result = {'device': device, 'passed': False, 'error': str(error)[:200]} results.append(result) print(f'QMD benchmark {device.upper()}: ' + ('passed' if result.get('passed') else 'failed') diff --git a/scripts/lib/qmd.sh b/scripts/lib/qmd.sh index d810a67..65834f4 100644 --- a/scripts/lib/qmd.sh +++ b/scripts/lib/qmd.sh @@ -1,4 +1,3 @@ -# scripts/lib/qmd.sh #!/usr/bin/env bash sync_qmd() { diff --git a/scripts/tests/test-qmd-lib.mjs b/scripts/tests/test-qmd-lib.mjs index bea60f1..997d2ce 100644 --- a/scripts/tests/test-qmd-lib.mjs +++ b/scripts/tests/test-qmd-lib.mjs @@ -5,7 +5,7 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import * as YAML from './fixtures/yaml-shim.mjs'; import { collectionName, deviceEnv, idleMinutes, qmdConfigPath, writePrivateJson, readJson, runQmdCli, MODELS } from '../../shared/qmd/qmd-lib.mjs'; -import { setModels, unsetModels } from '../../shared/qmd/qmd-config.mjs'; +import { setModels, unsetModels, previousModels, restoreModels } from '../../shared/qmd/qmd-config.mjs'; test('collection name is stable, prefixed, and path-safe', () => { const a = collectionName('/tmp/My Repo/ΓΌ'); @@ -75,3 +75,19 @@ test('runQmdCli runs the CLI entry of QMD_PACKAGE_DIR without a shell', () => { if (previous === undefined) delete process.env.QMD_PACKAGE_DIR; else process.env.QMD_PACKAGE_DIR = previous; } }); + +test('previousModels captures a user models block, null when absent', () => { + assert.deepEqual(previousModels('models:\n embed: mine\ncollections: {}\n', YAML), { embed: 'mine' }); + assert.equal(previousModels('collections: {}\n', YAML), null); + assert.equal(previousModels('', YAML), null); +}); + +test('restoreModels writes the previous block back and keeps comments and other keys', () => { + const original = '# keep\ncollections:\n a: {}\nmodels:\n embed: mine\n'; + const replaced = setModels(original, YAML); + const restored = restoreModels(replaced, YAML, { embed: 'mine' }); + assert.match(restored, /# keep/); + assert.match(restored, /embed: mine/); + assert.match(restored, /collections:/); + assert.doesNotMatch(restored, /Qwen/); +}); diff --git a/scripts/tests/test-qmd-setup.py b/scripts/tests/test-qmd-setup.py index 4a24143..fd743e0 100644 --- a/scripts/tests/test-qmd-setup.py +++ b/scripts/tests/test-qmd-setup.py @@ -1,10 +1,10 @@ -# scripts/tests/test-qmd-setup.py import importlib.util import json from pathlib import Path import shutil import tempfile import unittest +from unittest import mock ROOT = Path(__file__).resolve().parents[2] spec = importlib.util.spec_from_file_location('qmd_setup', ROOT / 'scripts/lib/qmd.py') @@ -13,15 +13,20 @@ spec.loader.exec_module(qmd) class FakeRunner: - def __init__(self, installed=False, gpu='cuda', gpu_result=None, cpu_result=None): + def __init__(self, installed=False, gpu='cuda', gpu_result=None, cpu_result=None, previous=None, fail_pull=False): self.calls, self.installed, self.gpu = [], installed, gpu + self.previous, self.fail_pull = previous, fail_pull self.results = {'gpu': gpu_result, 'cpu': cpu_result} - def __call__(self, arguments, env=None, timeout=None): + def __call__(self, arguments, env=None, timeout=None, check=True): self.calls.append(list(arguments)) joined = ' '.join(arguments) if arguments[:2] == ['node', '--version']: return 'v22.4.0' + if arguments[-1:] == ['set-models']: + return json.dumps({'previous': self.previous}) + chr(10) + if self.fail_pull and arguments[:2] == ['qmd', 'pull']: + raise ValueError('pull failed') if 'npm list' in joined: return json.dumps({'dependencies': {'@tobilu/qmd': {}} if self.installed else {}}) if 'npm install' in joined: @@ -126,6 +131,73 @@ class SetupTest(unittest.TestCase): self.assertEqual([r['kind'] for r in ledger['resources']], ['plugin'], 'plugin record stays so reconcile uninstalls qmd@qmd') self.assertTrue(self.state()['owned_package'], 'ledger-owned package is adopted') + def test_previous_models_are_restored_on_disable(self): + runner = FakeRunner(previous={'embed': 'mine'}) + setup = self.setup(runner) + setup.enable(['claude']) + self.assertEqual(self.state()['previous_models'], {'embed': 'mine'}) + setup.disable(['claude']) + self.assertIn(['node', str(self.home / '.my-ai-configuration/qmd/qmd-config.mjs'), 'restore-models', '{"embed": "mine"}'], runner.calls) + self.assertFalse(any(c[-1:] == ['unset-models'] for c in runner.calls)) + + def test_models_block_removed_when_none_existed(self): + runner = FakeRunner() + setup = self.setup(runner) + setup.enable(['claude']) + self.assertIsNone(self.state()['previous_models']) + setup.disable(['claude']) + self.assertTrue(any(c[-1:] == ['unset-models'] for c in runner.calls)) + self.assertFalse(any('restore-models' in c for c in runner.calls)) + + def ledger(self, base): + base.mkdir(parents=True, exist_ok=True) + (base / 'extensions.json').write_text(json.dumps({'version': 1, 'resources': [ + {'client': 'shared', 'name': 'QMD', 'kind': 'qmd', 'package': '@tobilu/qmd', 'clients': ['codex']}]})) + + def test_disable_without_state_uninstalls_adopted_package(self): + base = self.home / '.my-ai-configuration' + self.ledger(base) + runner = FakeRunner(installed=True) + self.setup(runner).disable(['claude']) + self.assertIn(['npm', 'uninstall', '--global', '@tobilu/qmd'], runner.calls) + self.assertFalse((base / 'qmd.json').exists()) + + def test_ownership_is_persisted_before_ledger_rewrite(self): + base = self.home / '.my-ai-configuration' + self.ledger(base) + real = qmd.atomic_json + + def failing(path, value): + if path.name == 'extensions.json': + raise OSError('disk full') + real(path, value) + with mock.patch.object(qmd, 'atomic_json', failing): + with self.assertRaises(OSError): + self.setup(FakeRunner(installed=True)).enable(['claude']) + self.assertTrue(self.state()['owned_package']) + self.assertEqual(len(json.loads((base / 'extensions.json').read_text())['resources']), 1) + + def test_symlinked_script_directory_is_refused(self): + base = self.home / '.my-ai-configuration' + base.mkdir(parents=True) + target = self.home / 'elsewhere' + target.mkdir() + try: + (base / 'qmd').symlink_to(target, target_is_directory=True) + except OSError: + self.skipTest('symlinks unavailable') + with self.assertRaises(ValueError): + self.setup(FakeRunner()).enable(['claude']) + self.assertEqual(list(target.iterdir()), []) + + def test_non_dict_benchmark_result_counts_as_failure(self): + runner = FakeRunner(gpu=False, cpu_result=['unexpected']) + self.assertFalse(self.setup(runner).benchmark()) + + def test_package_installed_parses_listing(self): + self.assertTrue(self.setup(FakeRunner(installed=True)).package_installed()) + self.assertFalse(self.setup(FakeRunner()).package_installed()) + def test_dry_run_changes_nothing(self): runner = FakeRunner() self.setup(runner, dry_run=True).enable(['claude']) diff --git a/shared/qmd/qmd-config.mjs b/shared/qmd/qmd-config.mjs index 4e836ea..f4f9d5d 100644 --- a/shared/qmd/qmd-config.mjs +++ b/shared/qmd/qmd-config.mjs @@ -12,6 +12,20 @@ export function setModels(text, YAML) { return document.toString(); } +export function previousModels(text, YAML) { + const document = YAML.parseDocument(text || ''); + if (document.errors.length) throw new Error(`index.yml cannot be parsed: ${document.errors[0].message}`); + return document.contents === null ? null : document.toJS().models ?? null; +} + +export function restoreModels(text, YAML, mapping) { + const document = YAML.parseDocument(text || ''); + if (document.errors.length) throw new Error(`index.yml cannot be parsed: ${document.errors[0].message}`); + if (document.contents === null) document.contents = document.createNode({}); + document.set('models', document.createNode(mapping)); + return document.toString(); +} + export function unsetModels(text, YAML) { const document = YAML.parseDocument(text || ''); if (document.errors.length) throw new Error(`index.yml cannot be parsed: ${document.errors[0].message}`); @@ -19,16 +33,26 @@ export function unsetModels(text, YAML) { return document.contents === null || document.contents.items?.length === 0 ? '' : document.toString(); } -async function main(action) { - if (!['set-models', 'unset-models'].includes(action)) throw new Error('usage: qmd-config.mjs set-models|unset-models'); +async function main(action, argument) { + if (!['set-models', 'unset-models', 'restore-models'].includes(action)) throw new Error('usage: qmd-config.mjs set-models|unset-models|restore-models '); const YAML = await import(pathToFileURL(join(qmdPackageDir(), 'node_modules', 'yaml', 'dist', 'index.js')).href); const path = qmdConfigPath(); const current = existsSync(path) ? readFileSync(path, 'utf8') : ''; - const next = action === 'set-models' ? setModels(current, YAML) : unsetModels(current, YAML); + let next; + if (action === 'set-models') { + const previous = previousModels(current, YAML); + // A block equal to ours is a leftover of an interrupted run, not the user's own. + console.log(JSON.stringify({ previous: JSON.stringify(previous) === JSON.stringify(MODELS) ? null : previous })); + next = setModels(current, YAML); + } else if (action === 'restore-models') { + const mapping = JSON.parse(argument ?? ''); + if (mapping === null || typeof mapping !== 'object' || Array.isArray(mapping)) throw new Error('restore-models needs a JSON object'); + next = restoreModels(current, YAML, mapping); + } else next = unsetModels(current, YAML); mkdirSync(dirname(path), { recursive: true }); writeFileSync(path, next); } if (import.meta.url === pathToFileURL(process.argv[1] || '').href) { - main(process.argv[2]).catch((error) => { console.error(`qmd-config: ${error.message}`); process.exit(1); }); + main(process.argv[2], process.argv[3]).catch((error) => { console.error(`qmd-config: ${error.message}`); process.exit(1); }); }