Update Astra: v2.80.19 - Refactoring Sidebar, LM Studio integration, and new tests

2026-05-08 23:14:47 +09:00
parent d083177d95
commit 5ffb472d22
28 changed files with 3125 additions and 1797 deletions
@@ -0,0 +1,80 @@
+/**
+ * Unit tests for findBrainFiles TTL cache.
+ */
+
+import * as fs from 'fs';
+import * as os from 'os';
+import * as path from 'path';
+import { findBrainFiles, invalidateBrainFilesCache } from '../src/utils';
+
+function makeBrain(): string {
+    const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'astra-brain-cache-'));
+    fs.writeFileSync(path.join(dir, 'a.md'), '# A');
+    fs.writeFileSync(path.join(dir, 'b.md'), '# B');
+    return dir;
+}
+
+describe('findBrainFiles TTL cache', () => {
+    let brain: string;
+
+    beforeEach(() => {
+        brain = makeBrain();
+        invalidateBrainFilesCache();
+    });
+
+    afterEach(() => {
+        try { fs.rmSync(brain, { recursive: true, force: true }); } catch { /* noop */ }
+    });
+
+    test('initial walk lists current files', () => {
+        const files = findBrainFiles(brain);
+        expect(files.length).toBe(2);
+        expect(files.some((f) => f.endsWith('a.md'))).toBe(true);
+    });
+
+    test('within TTL: cache returns stale list when files added', () => {
+        findBrainFiles(brain); // prime cache
+        fs.writeFileSync(path.join(brain, 'c.md'), '# C');
+        const files = findBrainFiles(brain);
+        // Cache hit returns the previous list
+        expect(files.length).toBe(2);
+    });
+
+    test('explicit invalidation forces fresh walk', () => {
+        findBrainFiles(brain); // prime
+        fs.writeFileSync(path.join(brain, 'c.md'), '# C');
+        invalidateBrainFilesCache(brain);
+        const files = findBrainFiles(brain);
+        expect(files.length).toBe(3);
+    });
+
+    test('invalidateBrainFilesCache() with no arg clears all entries', () => {
+        const a = makeBrain();
+        const b = makeBrain();
+        try {
+            findBrainFiles(a);
+            findBrainFiles(b);
+            fs.writeFileSync(path.join(a, 'extra.md'), 'x');
+            fs.writeFileSync(path.join(b, 'extra.md'), 'x');
+            invalidateBrainFilesCache();
+            expect(findBrainFiles(a).length).toBe(3);
+            expect(findBrainFiles(b).length).toBe(3);
+        } finally {
+            fs.rmSync(a, { recursive: true, force: true });
+            fs.rmSync(b, { recursive: true, force: true });
+        }
+    });
+
+    test('returned array is a copy — mutations do not poison cache', () => {
+        const first = findBrainFiles(brain);
+        first.length = 0; // mutate caller's copy
+        const second = findBrainFiles(brain);
+        expect(second.length).toBe(2);
+    });
+
+    test('non-existent directory returns empty list and does not throw', () => {
+        const fake = path.join(os.tmpdir(), 'astra-no-such-dir-' + Date.now());
+        const files = findBrainFiles(fake);
+        expect(files).toEqual([]);
+    });
+});
@@ -0,0 +1,310 @@
+/**
+ * Unit tests for ModelLifecycleManager.
+ *
+ * Strategy: inject mock ILMStudioClient and a simple in-memory IActivityTracker.
+ * No real LM Studio or SDK is touched — the manager file does not import the
+ * SDK directly (only types via `import type`).
+ */
+
+import {
+    ModelLifecycleManager,
+    LifecycleConfig,
+    LifecycleManagerDeps,
+} from '../src/lmstudio/lifecycleManager';
+import type { ILMStudioClient } from '../src/lmstudio/client';
+import type { IActivityTracker } from '../src/lmstudio/activityTracker';
+
+class FakeActivityTracker implements IActivityTracker {
+    private listeners: Array<(_: void) => any> = [];
+    public readonly onActivity = ((listener: (_: void) => any) => {
+        this.listeners.push(listener);
+        return { dispose: () => { this.listeners = this.listeners.filter(l => l !== listener); } };
+    }) as any;
+    bump(): void {
+        for (const l of this.listeners.slice()) l();
+    }
+}
+
+class FakeLMStudioClient implements ILMStudioClient {
+    public loadCalls: string[] = [];
+    public unloadCalls: string[] = [];
+    public listLoadedCalls = 0;
+    public failNextLoad: Error | null = null;
+    public failNextUnload: Error | null = null;
+    public loadDelayMs = 0;
+    public lastLoadSignal: AbortSignal | undefined;
+
+    setBaseUrl(_: string): void { /* noop */ }
+
+    async load(modelKey: string, signal?: AbortSignal): Promise<void> {
+        this.loadCalls.push(modelKey);
+        this.lastLoadSignal = signal;
+        if (this.loadDelayMs > 0) {
+            await new Promise<void>((resolve, reject) => {
+                const t = setTimeout(resolve, this.loadDelayMs);
+                if (signal) {
+                    const onAbort = () => { clearTimeout(t); reject(new Error('aborted')); };
+                    if (signal.aborted) onAbort();
+                    else signal.addEventListener('abort', onAbort);
+                }
+            });
+        }
+        if (this.failNextLoad) {
+            const err = this.failNextLoad;
+            this.failNextLoad = null;
+            throw err;
+        }
+    }
+
+    async unload(modelKey: string): Promise<void> {
+        this.unloadCalls.push(modelKey);
+        if (this.failNextUnload) {
+            const err = this.failNextUnload;
+            this.failNextUnload = null;
+            throw err;
+        }
+    }
+
+    async listLoaded(): Promise<string[]> {
+        this.listLoadedCalls++;
+        return [];
+    }
+
+    async isReachable(): Promise<boolean> {
+        return true;
+    }
+}
+
+function makeManager(overrides: Partial<LifecycleConfig> = {}, depOverrides: Partial<LifecycleManagerDeps> = {}) {
+    const client = new FakeLMStudioClient();
+    const activity = new FakeActivityTracker();
+    const config: LifecycleConfig = { idleTimeoutMs: 1000, autoLoadOnSelect: true, ...overrides };
+    const errors: string[] = [];
+    const manager = new ModelLifecycleManager({
+        client,
+        activity,
+        getConfig: () => config,
+        notifyError: (m) => errors.push(m),
+        switchDebounceMs: 0,
+        initialEngine: 'lmstudio',
+        ...depOverrides,
+    });
+    return { manager, client, activity, config, errors };
+}
+
+const flush = () => new Promise<void>((r) => setImmediate(r));
+
+describe('ModelLifecycleManager', () => {
+    beforeEach(() => {
+        jest.useFakeTimers({ doNotFake: ['setImmediate'] });
+    });
+
+    afterEach(async () => {
+        jest.useRealTimers();
+    });
+
+    test('loads model on selection and arms idle timer', async () => {
+        const { manager, client } = makeManager();
+        manager.onModelSelected('llama-3.2-3b');
+        await flush();
+        expect(client.loadCalls).toEqual(['llama-3.2-3b']);
+        expect(manager._getState()).toBe('loaded');
+        expect(manager._getCurrentModel()).toBe('llama-3.2-3b');
+        expect(manager._hasIdleTimer()).toBe(true);
+    });
+
+    test('idle timer triggers unload after timeout', async () => {
+        const { manager, client } = makeManager({ idleTimeoutMs: 1000 });
+        manager.onModelSelected('m1');
+        await flush();
+        expect(manager._getState()).toBe('loaded');
+
+        jest.advanceTimersByTime(1000);
+        await flush();
+        expect(client.unloadCalls).toEqual(['m1']);
+        expect(manager._getState()).toBe('idle');
+        expect(manager._getCurrentModel()).toBe(null);
+    });
+
+    test('idleTimeoutMs <= 0 disables auto-eject', async () => {
+        const { manager, client } = makeManager({ idleTimeoutMs: 0 });
+        manager.onModelSelected('m1');
+        await flush();
+        expect(manager._getState()).toBe('loaded');
+        expect(manager._hasIdleTimer()).toBe(false);
+
+        jest.advanceTimersByTime(1_000_000);
+        await flush();
+        expect(client.unloadCalls).toEqual([]);
+        expect(manager._getState()).toBe('loaded');
+    });
+
+    test('activity bump resets idle timer', async () => {
+        const { manager, client, activity } = makeManager({ idleTimeoutMs: 1000 });
+        manager.onModelSelected('m1');
+        await flush();
+
+        jest.advanceTimersByTime(900);
+        activity.bump();
+        jest.advanceTimersByTime(900); // total 1800ms but timer reset at 900
+        await flush();
+        expect(client.unloadCalls).toEqual([]);
+        expect(manager._getState()).toBe('loaded');
+
+        jest.advanceTimersByTime(200); // 200 + 900 since reset = ~1100 since reset
+        await flush();
+        expect(client.unloadCalls).toEqual(['m1']);
+    });
+
+    test('streamStart pauses idle timer; streamEnd resumes', async () => {
+        const { manager, client } = makeManager({ idleTimeoutMs: 500 });
+        manager.onModelSelected('m1');
+        await flush();
+        expect(manager._hasIdleTimer()).toBe(true);
+
+        manager.onStreamStart();
+        expect(manager._getState()).toBe('streaming');
+        expect(manager._hasIdleTimer()).toBe(false);
+
+        jest.advanceTimersByTime(10000);
+        await flush();
+        expect(client.unloadCalls).toEqual([]); // never ejected during stream
+
+        manager.onStreamEnd();
+        expect(manager._getState()).toBe('loaded');
+        expect(manager._hasIdleTimer()).toBe(true);
+
+        jest.advanceTimersByTime(500);
+        await flush();
+        expect(client.unloadCalls).toEqual(['m1']);
+    });
+
+    test('model switch unloads previous and loads next', async () => {
+        const { manager, client } = makeManager();
+        manager.onModelSelected('m1');
+        await flush();
+
+        manager.onModelSelected('m2');
+        await flush();
+        expect(client.unloadCalls).toEqual(['m1']);
+        expect(client.loadCalls).toEqual(['m1', 'm2']);
+        expect(manager._getCurrentModel()).toBe('m2');
+        expect(manager._getState()).toBe('loaded');
+    });
+
+    test('switch during streaming defers until streamEnd', async () => {
+        const { manager, client } = makeManager();
+        manager.onModelSelected('m1');
+        await flush();
+        manager.onStreamStart();
+
+        manager.onModelSelected('m2');
+        await flush();
+        // No switch yet — still in streaming with m1
+        expect(client.unloadCalls).toEqual([]);
+        expect(client.loadCalls).toEqual(['m1']);
+
+        manager.onStreamEnd();
+        await flush();
+        expect(client.unloadCalls).toEqual(['m1']);
+        expect(client.loadCalls).toEqual(['m1', 'm2']);
+        expect(manager._getCurrentModel()).toBe('m2');
+    });
+
+    test('load failure surfaces via notifyError and resets to idle', async () => {
+        const { manager, client, errors } = makeManager();
+        client.failNextLoad = new Error('LM Studio is not running');
+        manager.onModelSelected('m1');
+        await flush();
+        expect(manager._getState()).toBe('idle');
+        expect(manager._getCurrentModel()).toBe(null);
+        expect(errors.length).toBe(1);
+        expect(errors[0]).toContain('LM Studio is not running');
+    });
+
+    test('engine change to non-lmstudio clears state without unloading', async () => {
+        const { manager, client } = makeManager();
+        manager.onModelSelected('m1');
+        await flush();
+        expect(manager._getState()).toBe('loaded');
+
+        manager.setEngine('ollama');
+        expect(manager._getState()).toBe('idle');
+        expect(manager._getCurrentModel()).toBe(null);
+        expect(manager._hasIdleTimer()).toBe(false);
+        expect(client.unloadCalls).toEqual([]); // explicitly does not call unload
+    });
+
+    test('selection while engine is non-lmstudio is a no-op', async () => {
+        const { manager, client } = makeManager({}, { initialEngine: 'ollama' });
+        manager.onModelSelected('m1');
+        await flush();
+        expect(client.loadCalls).toEqual([]);
+        expect(manager._getState()).toBe('idle');
+    });
+
+    test('autoLoadOnSelect=false skips loading', async () => {
+        const { manager, client } = makeManager({ autoLoadOnSelect: false });
+        manager.onModelSelected('m1');
+        await flush();
+        expect(client.loadCalls).toEqual([]);
+        expect(manager._getState()).toBe('idle');
+    });
+
+    test('rapid switch debounce: only the last selection wins', async () => {
+        // Re-create manager with non-zero debounce
+        const client = new FakeLMStudioClient();
+        const activity = new FakeActivityTracker();
+        const config: LifecycleConfig = { idleTimeoutMs: 1000, autoLoadOnSelect: true };
+        const manager = new ModelLifecycleManager({
+            client,
+            activity,
+            getConfig: () => config,
+            switchDebounceMs: 300,
+            initialEngine: 'lmstudio',
+        });
+        manager.onModelSelected('m1');
+        manager.onModelSelected('m2');
+        manager.onModelSelected('m3');
+        // Before debounce expires no load fires
+        expect(client.loadCalls).toEqual([]);
+        jest.advanceTimersByTime(300);
+        await flush();
+        expect(client.loadCalls).toEqual(['m3']);
+    });
+
+    test('disposeAndUnload ejects loaded model', async () => {
+        const { manager, client } = makeManager();
+        manager.onModelSelected('m1');
+        await flush();
+        await manager.disposeAndUnload(2000);
+        expect(client.unloadCalls).toEqual(['m1']);
+        expect(manager._getState()).toBe('idle');
+    });
+
+    test('disposeAndUnload while idle is a no-op', async () => {
+        const { manager, client } = makeManager();
+        await manager.disposeAndUnload(500);
+        expect(client.unloadCalls).toEqual([]);
+    });
+
+    test('selecting same model does not re-load but refreshes timer', async () => {
+        const { manager, client } = makeManager({ idleTimeoutMs: 1000 });
+        manager.onModelSelected('m1');
+        await flush();
+        expect(client.loadCalls.length).toBe(1);
+
+        jest.advanceTimersByTime(800);
+        manager.onModelSelected('m1');
+        await flush();
+        expect(client.loadCalls.length).toBe(1); // no reload
+        // Timer was reset; should survive past original 1000ms
+        jest.advanceTimersByTime(700);
+        await flush();
+        expect(client.unloadCalls).toEqual([]);
+
+        jest.advanceTimersByTime(400);
+        await flush();
+        expect(client.unloadCalls).toEqual(['m1']);
+    });
+});
@@ -7,10 +7,28 @@ const config = {
  }
 };

+class EventEmitter {
+  constructor() {
+    this._listeners = [];
+    this.event = (listener) => {
+      this._listeners.push(listener);
+      return { dispose: () => { this._listeners = this._listeners.filter(l => l !== listener); } };
+    };
+  }
+  fire(value) {
+    for (const l of this._listeners.slice()) l(value);
+  }
+  dispose() {
+    this._listeners = [];
+  }
+}
+
 module.exports = {
+  EventEmitter,
  workspace: {
    workspaceFolders: [],
    getConfiguration: () => config,
+    onDidChangeConfiguration: () => ({ dispose: () => {} }),
    fs: {
      writeFile: async () => {},
      delete: async () => {}