Update Astra: v2.80.19 - Refactoring Sidebar, LM Studio integration, and new tests
This commit is contained in:
g1nation
@@ -0,0 +1,19 @@
|
||||
import * as vscode from 'vscode';
|
||||
|
||||
export interface IActivityTracker {
|
||||
onActivity: vscode.Event<void>;
|
||||
bump(): void;
|
||||
}
|
||||
|
||||
export class ActivityTracker implements IActivityTracker {
|
||||
private readonly _emitter = new vscode.EventEmitter<void>();
|
||||
public readonly onActivity = this._emitter.event;
|
||||
|
||||
bump(): void {
|
||||
this._emitter.fire();
|
||||
}
|
||||
|
||||
dispose(): void {
|
||||
this._emitter.dispose();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
import { LMStudioClient as SDKClient } from '@lmstudio/sdk';
|
||||
import { logError, logInfo } from '../utils';
|
||||
|
||||
export interface ILMStudioClient {
|
||||
load(modelKey: string, signal?: AbortSignal): Promise<void>;
|
||||
unload(modelKey: string): Promise<void>;
|
||||
listLoaded(): Promise<string[]>;
|
||||
isReachable(): Promise<boolean>;
|
||||
setBaseUrl(httpBaseUrl: string): void;
|
||||
}
|
||||
|
||||
export class LMStudioLifecycleError extends Error {
|
||||
constructor(message: string, public readonly cause?: unknown) {
|
||||
super(message);
|
||||
this.name = 'LMStudioLifecycleError';
|
||||
}
|
||||
}
|
||||
|
||||
export function httpToWebSocketUrl(httpBaseUrl: string): string | undefined {
|
||||
const trimmed = (httpBaseUrl || '').trim();
|
||||
if (!trimmed) return undefined;
|
||||
try {
|
||||
const url = new URL(trimmed);
|
||||
if (url.protocol === 'http:') url.protocol = 'ws:';
|
||||
else if (url.protocol === 'https:') url.protocol = 'wss:';
|
||||
else if (url.protocol !== 'ws:' && url.protocol !== 'wss:') return undefined;
|
||||
if (url.pathname.endsWith('/v1')) url.pathname = url.pathname.slice(0, -3);
|
||||
if (url.pathname.endsWith('/api')) url.pathname = url.pathname.slice(0, -4);
|
||||
const out = url.toString().replace(/\/+$/, '');
|
||||
return out;
|
||||
} catch {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
export class LMStudioClient implements ILMStudioClient {
|
||||
private _sdk: SDKClient | undefined;
|
||||
private _wsUrl: string | undefined;
|
||||
|
||||
constructor(httpBaseUrl: string) {
|
||||
this.setBaseUrl(httpBaseUrl);
|
||||
}
|
||||
|
||||
setBaseUrl(httpBaseUrl: string): void {
|
||||
const ws = httpToWebSocketUrl(httpBaseUrl);
|
||||
if (ws !== this._wsUrl) {
|
||||
this._wsUrl = ws;
|
||||
this._sdk = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
private getSdk(): SDKClient {
|
||||
if (!this._sdk) {
|
||||
this._sdk = new SDKClient(this._wsUrl ? { baseUrl: this._wsUrl } : {});
|
||||
}
|
||||
return this._sdk;
|
||||
}
|
||||
|
||||
async load(modelKey: string, signal?: AbortSignal): Promise<void> {
|
||||
try {
|
||||
await this.getSdk().llm.load(modelKey, signal ? { signal } : undefined);
|
||||
logInfo('LM Studio model loaded.', { modelKey });
|
||||
} catch (e: any) {
|
||||
const msg = e?.message ?? String(e);
|
||||
throw new LMStudioLifecycleError(`Failed to load LM Studio model "${modelKey}": ${msg}`, e);
|
||||
}
|
||||
}
|
||||
|
||||
async unload(modelKey: string): Promise<void> {
|
||||
try {
|
||||
await this.getSdk().llm.unload(modelKey);
|
||||
logInfo('LM Studio model unloaded.', { modelKey });
|
||||
} catch (e: any) {
|
||||
const msg = e?.message ?? String(e);
|
||||
throw new LMStudioLifecycleError(`Failed to unload LM Studio model "${modelKey}": ${msg}`, e);
|
||||
}
|
||||
}
|
||||
|
||||
async listLoaded(): Promise<string[]> {
|
||||
try {
|
||||
const items: any[] = await this.getSdk().llm.listLoaded();
|
||||
return items
|
||||
.map((m) => m?.identifier ?? m?.modelKey ?? m?.path ?? null)
|
||||
.filter((id): id is string => typeof id === 'string' && id.length > 0);
|
||||
} catch (e: any) {
|
||||
const msg = e?.message ?? String(e);
|
||||
throw new LMStudioLifecycleError(`Failed to list loaded LM Studio models: ${msg}`, e);
|
||||
}
|
||||
}
|
||||
|
||||
async isReachable(): Promise<boolean> {
|
||||
try {
|
||||
await this.getSdk().llm.listLoaded();
|
||||
return true;
|
||||
} catch (e: any) {
|
||||
logError('LM Studio not reachable.', { error: e?.message ?? String(e) });
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
import type { ILMStudioClient } from './client';
|
||||
import type { IActivityTracker } from './activityTracker';
|
||||
import type { EngineKind } from '../utils';
|
||||
import { logError, logInfo } from '../utils';
|
||||
|
||||
export type LifecycleState = 'idle' | 'loading' | 'loaded' | 'streaming' | 'unloading';
|
||||
|
||||
export interface LifecycleConfig {
|
||||
idleTimeoutMs: number;
|
||||
autoLoadOnSelect: boolean;
|
||||
}
|
||||
|
||||
export interface LifecycleManagerDeps {
|
||||
client: ILMStudioClient;
|
||||
activity: IActivityTracker;
|
||||
getConfig: () => LifecycleConfig;
|
||||
notifyError?: (msg: string) => void;
|
||||
/** Debounce window for rapid model switches. Default 300ms. Use 0 in tests for synchronous behavior. */
|
||||
switchDebounceMs?: number;
|
||||
/** Initial engine. Default 'lmstudio'. */
|
||||
initialEngine?: EngineKind;
|
||||
}
|
||||
|
||||
export class ModelLifecycleManager {
|
||||
private state: LifecycleState = 'idle';
|
||||
private currentModel: string | null = null;
|
||||
private pendingModel: string | null = null;
|
||||
private engine: EngineKind;
|
||||
|
||||
private idleTimer: ReturnType<typeof setTimeout> | undefined;
|
||||
private switchDebounce: ReturnType<typeof setTimeout> | undefined;
|
||||
private loadAbort: AbortController | undefined;
|
||||
|
||||
private readonly activitySub: { dispose(): void };
|
||||
private disposed = false;
|
||||
|
||||
constructor(private readonly deps: LifecycleManagerDeps) {
|
||||
this.engine = deps.initialEngine ?? 'lmstudio';
|
||||
this.activitySub = deps.activity.onActivity(() => this.onActivity());
|
||||
}
|
||||
|
||||
setEngine(engine: EngineKind): void {
|
||||
if (engine === this.engine) return;
|
||||
const wasLmStudio = this.engine === 'lmstudio';
|
||||
this.engine = engine;
|
||||
if (wasLmStudio && engine !== 'lmstudio') {
|
||||
this.clearIdleTimer();
|
||||
this.cancelPendingSwitch();
|
||||
this.cancelLoad();
|
||||
this.state = 'idle';
|
||||
this.currentModel = null;
|
||||
this.pendingModel = null;
|
||||
}
|
||||
}
|
||||
|
||||
onModelSelected(modelKey: string): void {
|
||||
if (this.disposed) return;
|
||||
if (this.engine !== 'lmstudio') return;
|
||||
if (!this.deps.getConfig().autoLoadOnSelect) return;
|
||||
const trimmed = (modelKey || '').trim();
|
||||
if (!trimmed) return;
|
||||
|
||||
// Mid-stream: queue the latest selection, apply on streamEnd.
|
||||
if (this.state === 'streaming') {
|
||||
this.pendingModel = trimmed;
|
||||
return;
|
||||
}
|
||||
|
||||
// Same model already in flight or active — keep timer fresh, no reload.
|
||||
if ((this.state === 'loaded' || this.state === 'loading') && this.currentModel === trimmed) {
|
||||
if (this.state === 'loaded') this.resetIdleTimer();
|
||||
return;
|
||||
}
|
||||
|
||||
this.cancelPendingSwitch();
|
||||
const delay = this.deps.switchDebounceMs ?? 300;
|
||||
if (delay <= 0) {
|
||||
void this.doSwitch(trimmed);
|
||||
return;
|
||||
}
|
||||
this.switchDebounce = setTimeout(() => {
|
||||
this.switchDebounce = undefined;
|
||||
void this.doSwitch(trimmed);
|
||||
}, delay);
|
||||
}
|
||||
|
||||
onStreamStart(): void {
|
||||
if (this.disposed) return;
|
||||
if (this.engine !== 'lmstudio') return;
|
||||
this.clearIdleTimer();
|
||||
if (this.state === 'loaded') this.state = 'streaming';
|
||||
}
|
||||
|
||||
onStreamEnd(): void {
|
||||
if (this.disposed) return;
|
||||
if (this.engine !== 'lmstudio') return;
|
||||
if (this.state === 'streaming') {
|
||||
this.state = 'loaded';
|
||||
if (this.pendingModel && this.pendingModel !== this.currentModel) {
|
||||
const next = this.pendingModel;
|
||||
this.pendingModel = null;
|
||||
void this.doSwitch(next);
|
||||
} else {
|
||||
this.pendingModel = null;
|
||||
this.resetIdleTimer();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Best-effort eject before extension shutdown. Bounded by timeoutMs. */
|
||||
async disposeAndUnload(timeoutMs: number = 2000): Promise<void> {
|
||||
if (this.disposed) return;
|
||||
this.disposed = true;
|
||||
this.clearIdleTimer();
|
||||
this.cancelPendingSwitch();
|
||||
this.cancelLoad();
|
||||
this.activitySub.dispose();
|
||||
|
||||
const shouldUnload =
|
||||
this.engine === 'lmstudio' &&
|
||||
(this.state === 'loaded' || this.state === 'streaming') &&
|
||||
this.currentModel !== null;
|
||||
if (!shouldUnload) {
|
||||
this.state = 'idle';
|
||||
this.currentModel = null;
|
||||
return;
|
||||
}
|
||||
|
||||
const target = this.currentModel as string;
|
||||
this.state = 'unloading';
|
||||
try {
|
||||
await Promise.race([
|
||||
this.deps.client.unload(target),
|
||||
new Promise<void>((_, reject) =>
|
||||
setTimeout(() => reject(new Error(`unload timed out after ${timeoutMs}ms`)), timeoutMs)
|
||||
),
|
||||
]);
|
||||
} catch (e: any) {
|
||||
logError('LM Studio unload during dispose failed.', { model: target, error: e?.message ?? String(e) });
|
||||
}
|
||||
this.state = 'idle';
|
||||
this.currentModel = null;
|
||||
}
|
||||
|
||||
/** vscode.Disposable shape — fire and forget. */
|
||||
dispose(): void {
|
||||
void this.disposeAndUnload();
|
||||
}
|
||||
|
||||
// Test/inspection helpers
|
||||
public _getState(): LifecycleState { return this.state; }
|
||||
public _getCurrentModel(): string | null { return this.currentModel; }
|
||||
public _hasIdleTimer(): boolean { return this.idleTimer !== undefined; }
|
||||
|
||||
// ---------- internals ----------
|
||||
|
||||
private onActivity(): void {
|
||||
if (this.disposed) return;
|
||||
if (this.engine !== 'lmstudio') return;
|
||||
if (this.state !== 'loaded') return;
|
||||
this.resetIdleTimer();
|
||||
}
|
||||
|
||||
private clearIdleTimer(): void {
|
||||
if (this.idleTimer) {
|
||||
clearTimeout(this.idleTimer);
|
||||
this.idleTimer = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
private cancelPendingSwitch(): void {
|
||||
if (this.switchDebounce) {
|
||||
clearTimeout(this.switchDebounce);
|
||||
this.switchDebounce = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
private resetIdleTimer(): void {
|
||||
this.clearIdleTimer();
|
||||
const ms = this.deps.getConfig().idleTimeoutMs;
|
||||
if (!Number.isFinite(ms) || ms <= 0) return;
|
||||
this.idleTimer = setTimeout(() => {
|
||||
this.idleTimer = undefined;
|
||||
void this.doIdleEject();
|
||||
}, ms);
|
||||
}
|
||||
|
||||
private async doIdleEject(): Promise<void> {
|
||||
if (this.state !== 'loaded' || !this.currentModel) return;
|
||||
const target = this.currentModel;
|
||||
this.state = 'unloading';
|
||||
try {
|
||||
await this.deps.client.unload(target);
|
||||
logInfo('LM Studio model auto-ejected after idle.', { model: target });
|
||||
} catch (e: any) {
|
||||
logError('LM Studio auto-eject failed.', { model: target, error: e?.message ?? String(e) });
|
||||
this.deps.notifyError?.(`LM Studio auto-eject failed: ${e?.message ?? e}`);
|
||||
}
|
||||
this.state = 'idle';
|
||||
this.currentModel = null;
|
||||
}
|
||||
|
||||
private cancelLoad(): void {
|
||||
if (this.loadAbort) {
|
||||
try { this.loadAbort.abort(); } catch { /* noop */ }
|
||||
this.loadAbort = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
private async doSwitch(modelKey: string): Promise<void> {
|
||||
if (this.disposed) return;
|
||||
if (this.engine !== 'lmstudio') return;
|
||||
|
||||
this.cancelLoad();
|
||||
this.clearIdleTimer();
|
||||
|
||||
if (this.state === 'loaded' && this.currentModel && this.currentModel !== modelKey) {
|
||||
const prev = this.currentModel;
|
||||
this.state = 'unloading';
|
||||
try {
|
||||
await this.deps.client.unload(prev);
|
||||
} catch (e: any) {
|
||||
logError('LM Studio unload before switch failed.', { prev, error: e?.message ?? String(e) });
|
||||
}
|
||||
this.currentModel = null;
|
||||
}
|
||||
|
||||
this.state = 'loading';
|
||||
this.currentModel = modelKey;
|
||||
const ac = new AbortController();
|
||||
this.loadAbort = ac;
|
||||
try {
|
||||
await this.deps.client.load(modelKey, ac.signal);
|
||||
if (this.loadAbort !== ac) return; // superseded by a newer switch
|
||||
this.loadAbort = undefined;
|
||||
this.state = 'loaded';
|
||||
this.resetIdleTimer();
|
||||
} catch (e: any) {
|
||||
if (ac.signal.aborted) return; // superseded — newer switch owns state
|
||||
logError('LM Studio model load failed.', { model: modelKey, error: e?.message ?? String(e) });
|
||||
this.deps.notifyError?.(`LM Studio load failed: ${e?.message ?? e}`);
|
||||
if (this.loadAbort === ac) this.loadAbort = undefined;
|
||||
this.state = 'idle';
|
||||
this.currentModel = null;
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user