diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts index a171b7b68..3b0a9d396 100644 --- a/packages/backend/src/managers/modelsManager.spec.ts +++ b/packages/backend/src/managers/modelsManager.spec.ts @@ -225,7 +225,7 @@ test('getModelsInfo should get models in local directory', async () => { modelHandlerRegistry, ); modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir)); - manager.init(); + await manager.init(); await manager.loadLocalModels(); expect(manager.getModelsInfo()).toEqual([ { @@ -277,7 +277,7 @@ test('getModelsInfo should return an empty array if the models folder does not e modelHandlerRegistry, ); modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir)); - manager.init(); + await manager.init(); await manager.getLocalModelsFromDisk(); expect(manager.getModelsInfo()).toEqual([]); if (process.platform === 'win32') { @@ -318,7 +318,7 @@ test('getLocalModelsFromDisk should return undefined Date and size when stat fai modelHandlerRegistry, ); modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir)); - manager.init(); + await manager.init(); await manager.loadLocalModels(); expect(manager.getModelsInfo()).toEqual([ { @@ -377,7 +377,7 @@ test('getLocalModelsFromDisk should skip folders containing tmp files', async () modelHandlerRegistry, ); modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir)); - manager.init(); + await manager.init(); await manager.loadLocalModels(); expect(manager.getModelsInfo()).toEqual([ { @@ -417,7 +417,7 @@ test('loadLocalModels should post a message with the message on disk and on cata modelHandlerRegistry, ); modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir)); - manager.init(); + await manager.init(); await manager.loadLocalModels(); expect(rpcExtensionMock.fire).toHaveBeenNthCalledWith(2, MSG_NEW_MODELS_STATE, [ { @@ -464,7 +464,7 @@ test('deleteModel deletes the model folder', async () => { modelHandlerRegistry, ); modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir)); - manager.init(); + await manager.init(); await manager.loadLocalModels(); await manager.deleteModel('model-id-1'); // check that the model's folder is removed from disk @@ -525,7 +525,7 @@ describe('deleting models', () => { modelHandlerRegistry, ); modelHandlerRegistry.register(new URLModelHandler(manager, modelsDir)); - manager.init(); + await manager.init(); await manager.loadLocalModels(); await manager.deleteModel('model-id-1'); // check that the model's folder is removed from disk @@ -899,7 +899,7 @@ describe('getModelMetadata', () => { modelHandlerRegistry, ); - manager.init(); + await manager.init(); const fakeMetadata: Record = { hello: 'world', @@ -939,7 +939,7 @@ describe('getModelMetadata', () => { modelHandlerRegistry, ); - manager.init(); + await manager.init(); const fakeMetadata: Record = { hello: 'world', @@ -995,7 +995,7 @@ describe('uploadModelToPodmanMachine', () => { modelHandlerRegistry, ); - manager.init(); + await manager.init(); const result = await manager.uploadModelToPodmanMachine(connectionMock, modelMock); expect(result).toBe('uploader-result'); expect(performMock).toHaveBeenCalledWith(modelMock.id); @@ -1028,7 +1028,7 @@ describe('uploadModelToPodmanMachine', () => { modelHandlerRegistry, ); - manager.init(); + await manager.init(); await manager.uploadModelToPodmanMachine(connectionMock, modelMock); expect(Uploader).not.toHaveBeenCalled(); }); diff --git a/packages/backend/src/managers/modelsManager.ts b/packages/backend/src/managers/modelsManager.ts index 09a26c0cf..0633486e7 100644 --- a/packages/backend/src/managers/modelsManager.ts +++ b/packages/backend/src/managers/modelsManager.ts @@ -63,7 +63,7 @@ export class ModelsManager implements Disposable { this.modelHandlerRegistry.getAll().forEach(handler => handler.onUpdate(this.loadLocalModels)); } - init(): void { + async init(): Promise { const disposable = this.catalogManager.onUpdate(() => { this.loadLocalModels().catch((err: unknown) => { console.error(`Something went wrong when loading local models`, err); @@ -71,9 +71,11 @@ export class ModelsManager implements Disposable { }); this.#disposables.push(disposable); - this.loadLocalModels().catch((err: unknown) => { + try { + await this.loadLocalModels(); + } catch (err: unknown) { console.error('Something went wrong while trying to load local models', err); - }); + } } dispose(): void { diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts index 0f6f051cd..c54d2c9c5 100644 --- a/packages/backend/src/studio.ts +++ b/packages/backend/src/studio.ts @@ -251,7 +251,7 @@ export class Studio { const hfModelHandler = new HuggingFaceModelHandler(this.#modelsManager); this.#extensionContext.subscriptions.push(hfModelHandler); this.#extensionContext.subscriptions.push(modelHandlerRegistry.register(hfModelHandler)); - this.#modelsManager.init(); + await this.#modelsManager.init(); this.#extensionContext.subscriptions.push(this.#modelsManager); /**