connectai/tests/lmStudioLifecycle.test.ts

/**
 * 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;
    }

    async listLoadedCached(): Promise<string[]> {
        return [];
    }

    async listDownloaded(): Promise<string[]> {
        return [];
    }

    async getModelHandle(_modelKey: string): Promise<any> {
        return {};
    }
}

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']);
    });
});