diff --git a/packages/backend/src/managers/playground.spec.ts b/packages/backend/src/managers/playground.spec.ts index 3ff7df873..21cc07bc0 100644 --- a/packages/backend/src/managers/playground.spec.ts +++ b/packages/backend/src/managers/playground.spec.ts @@ -19,6 +19,7 @@ import { beforeEach, expect, test, vi } from 'vitest'; import { PlayGroundManager } from './playground'; import type { ImageInfo, Webview } from '@podman-desktop/api'; +import type { ContainerRegistry } from '../registries/ContainerRegistry'; const mocks = vi.hoisted(() => ({ postMessage: vi.fn(), @@ -27,6 +28,7 @@ const mocks = vi.hoisted(() => ({ createContainer: vi.fn(), stopContainer: vi.fn(), getFreePort: vi.fn(), + containerRegistrySubscribeMock: vi.fn(), })); vi.mock('@podman-desktop/api', async () => { @@ -42,6 +44,10 @@ vi.mock('@podman-desktop/api', async () => { }; }); +const containerRegistryMock = { + subscribe: mocks.containerRegistrySubscribeMock, +} as unknown as ContainerRegistry; + vi.mock('../utils/ports', async () => { return { getFreePort: mocks.getFreePort, @@ -53,9 +59,12 @@ let manager: PlayGroundManager; beforeEach(() => { vi.resetAllMocks(); - manager = new PlayGroundManager({ - postMessage: mocks.postMessage, - } as unknown as Webview); + manager = new PlayGroundManager( + { + postMessage: mocks.postMessage, + } as unknown as Webview, + containerRegistryMock, + ); }); test('startPlayground should fail if no provider', async () => { diff --git a/packages/backend/src/managers/playground.ts b/packages/backend/src/managers/playground.ts index 67e02d261..6c7de8f8e 100644 --- a/packages/backend/src/managers/playground.ts +++ b/packages/backend/src/managers/playground.ts @@ -32,6 +32,7 @@ import { getFreePort } from '../utils/ports'; import type { QueryState } from '@shared/src/models/IPlaygroundQueryState'; import { MSG_NEW_PLAYGROUND_QUERIES_STATE, MSG_PLAYGROUNDS_STATE_UPDATE } from '@shared/Messages'; import type { PlaygroundState, PlaygroundStatus } from '@shared/src/models/IPlaygroundState'; +import type { ContainerRegistry } from '../registries/ContainerRegistry'; // TODO: this should not be hardcoded const PLAYGROUND_IMAGE = 'quay.io/bootsy/playground:v0'; @@ -51,7 +52,10 @@ export class PlayGroundManager { private playgrounds: Map; private queries: Map; - constructor(private webview: Webview) { + constructor( + private webview: Webview, + private containerRegistry: ContainerRegistry, + ) { this.playgrounds = new Map(); this.queries = new Map(); } @@ -145,6 +149,21 @@ export class PlayGroundManager { Cmd: ['--models-path', '/models', '--context-size', '700', '--threads', '4'], }); + const disposable = this.containerRegistry.subscribe(result.id, (status: string) => { + switch (status) { + case 'remove': + case 'die': + case 'cleanup': + // Update the playground state accordingly + this.updatePlaygroundState(modelId, { + status: 'none', + modelId, + }); + disposable.dispose(); + break; + } + }); + this.updatePlaygroundState(modelId, { container: { containerId: result.id, diff --git a/packages/backend/src/registries/ContainerRegistry.spec.ts b/packages/backend/src/registries/ContainerRegistry.spec.ts new file mode 100644 index 000000000..3e0eaac82 --- /dev/null +++ b/packages/backend/src/registries/ContainerRegistry.spec.ts @@ -0,0 +1,123 @@ +/********************************************************************** + * Copyright (C) 2024 Red Hat, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + ***********************************************************************/ +import { expect, test, vi } from 'vitest'; +import { ContainerRegistry } from './ContainerRegistry'; +import type { ContainerJSONEvent } from '@podman-desktop/api'; + +const mocks = vi.hoisted(() => ({ + onEventMock: vi.fn(), + DisposableCreateMock: vi.fn(), +})); + +vi.mock('@podman-desktop/api', async () => { + return { + Disposable: { + create: mocks.DisposableCreateMock, + }, + containerEngine: { + onEvent: mocks.onEventMock, + }, + }; +}); + +test('ContainerRegistry init', () => { + const registry = new ContainerRegistry(); + registry.init(); + + expect(mocks.onEventMock).toHaveBeenCalledOnce(); +}); + +test('ContainerRegistry subscribe', () => { + // Get the callback created by the ContainerRegistry + let callback: (event: ContainerJSONEvent) => void; + mocks.onEventMock.mockImplementation((method: (event: ContainerJSONEvent) => void) => { + callback = method; + }); + + // Create the ContainerRegistry and init + const registry = new ContainerRegistry(); + registry.init(); + + // Let's create a dummy subscriber + let subscribedStatus: undefined | string = undefined; + registry.subscribe('random', (status: string) => { + subscribedStatus = status; + }); + + // Generate a fake event + callback({ + status: 'die', + id: 'random', + type: 'container', + }); + + expect(subscribedStatus).toBe('die'); + expect(mocks.DisposableCreateMock).toHaveBeenCalledOnce(); +}); + +test('ContainerRegistry unsubscribe all if container remove', () => { + // Get the callback created by the ContainerRegistry + let callback: (event: ContainerJSONEvent) => void; + mocks.onEventMock.mockImplementation((method: (event: ContainerJSONEvent) => void) => { + callback = method; + }); + + // Create the ContainerRegistry and init + const registry = new ContainerRegistry(); + registry.init(); + + // Let's create a dummy subscriber + const subscribeMock = vi.fn(); + registry.subscribe('random', subscribeMock); + + // Generate a remove event + callback({ status: 'remove', id: 'random', type: 'container' }); + + // Call it a second time + callback({ status: 'remove', id: 'random', type: 'container' }); + + // Our subscriber should only have been called once, the first, after it should have been removed. + expect(subscribeMock).toHaveBeenCalledOnce(); +}); + +test('ContainerRegistry subscriber disposed should not be called', () => { + // Get the callback created by the ContainerRegistry + let callback: (event: ContainerJSONEvent) => void; + mocks.onEventMock.mockImplementation((method: (event: ContainerJSONEvent) => void) => { + callback = method; + }); + + mocks.DisposableCreateMock.mockImplementation(callback => ({ + dispose: () => callback(), + })); + + // Create the ContainerRegistry and init + const registry = new ContainerRegistry(); + registry.init(); + + // Let's create a dummy subscriber + const subscribeMock = vi.fn(); + const disposable = registry.subscribe('random', subscribeMock); + disposable.dispose(); + + // Generate a random event + callback({ status: 'die', id: 'random', type: 'container' }); + + // never should have been called + expect(subscribeMock).toHaveBeenCalledTimes(0); +}); diff --git a/packages/backend/src/registries/ContainerRegistry.ts b/packages/backend/src/registries/ContainerRegistry.ts new file mode 100644 index 000000000..29aa7beb7 --- /dev/null +++ b/packages/backend/src/registries/ContainerRegistry.ts @@ -0,0 +1,62 @@ +/********************************************************************** + * Copyright (C) 2024 Red Hat, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + ***********************************************************************/ +import * as podmanDesktopApi from '@podman-desktop/api'; + +export type Subscriber = { + id: number; + callback: (status: string) => void; +}; + +export class ContainerRegistry { + private count: number = 0; + private subscribers: Map = new Map(); + + init(): podmanDesktopApi.Disposable { + return podmanDesktopApi.containerEngine.onEvent(event => { + if (this.subscribers.has(event.id)) { + this.subscribers.get(event.id).forEach(subscriber => subscriber.callback(event.status)); + + // If the event type is remove, we dispose all subscribers for the specific containers + if (event.status === 'remove') { + this.subscribers.delete(event.id); + } + } + }); + } + + subscribe(containerId: string, callback: (status: string) => void): podmanDesktopApi.Disposable { + const existing: Subscriber[] = this.subscribers.has(containerId) ? this.subscribers.get(containerId) : []; + const subscriberId = ++this.count; + this.subscribers.set(containerId, [ + { + id: subscriberId, + callback: callback, + }, + ...existing, + ]); + + return podmanDesktopApi.Disposable.create(() => { + if (!this.subscribers.has(containerId)) return; + + this.subscribers.set( + containerId, + this.subscribers.get(containerId).filter(subscriber => subscriber.id !== subscriberId), + ); + }); + } +} diff --git a/packages/backend/src/studio.spec.ts b/packages/backend/src/studio.spec.ts index f15bec5d8..bf7bb8e2a 100644 --- a/packages/backend/src/studio.spec.ts +++ b/packages/backend/src/studio.spec.ts @@ -46,6 +46,9 @@ vi.mock('@podman-desktop/api', async () => { }, }), }, + containerEngine: { + onEvent: vi.fn(), + }, }; }); diff --git a/packages/backend/src/studio.ts b/packages/backend/src/studio.ts index a8aa432c9..9b4563706 100644 --- a/packages/backend/src/studio.ts +++ b/packages/backend/src/studio.ts @@ -30,6 +30,7 @@ import { ModelsManager } from './managers/modelsManager'; import path from 'node:path'; import os from 'os'; import fs from 'node:fs'; +import { ContainerRegistry } from './registries/ContainerRegistry'; // TODO: Need to be configured export const AI_STUDIO_FOLDER = path.join('podman-desktop', 'ai-studio'); @@ -94,6 +95,10 @@ export class Studio { this.#panel.webview.html = indexHtml; + // Creating container registry + const containerRegistry = new ContainerRegistry(); + this.#extensionContext.subscriptions.push(containerRegistry.init()); + // Let's create the api that the front will be able to call const appUserDirectory = path.join(os.homedir(), AI_STUDIO_FOLDER); @@ -101,7 +106,7 @@ export class Studio { const gitManager = new GitManager(); const taskRegistry = new TaskRegistry(); const recipeStatusRegistry = new RecipeStatusRegistry(taskRegistry, this.#panel.webview); - this.playgroundManager = new PlayGroundManager(this.#panel.webview); + this.playgroundManager = new PlayGroundManager(this.#panel.webview, containerRegistry); // Create catalog manager, responsible for loading the catalog files and watching for changes this.catalogManager = new CatalogManager(appUserDirectory, this.#panel.webview); this.modelsManager = new ModelsManager(appUserDirectory, this.#panel.webview, this.catalogManager);