From 206a57564f95de60eef1ffbec4ef29c9aa089adb Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Tue, 30 Jan 2024 14:50:12 +0100 Subject: [PATCH 1/7] feat: listen to podman desktop container api events Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- packages/backend/src/managers/playground.ts | 21 ++++++- .../src/registries/ContainerRegistry.ts | 62 +++++++++++++++++++ packages/backend/src/studio.ts | 7 ++- 3 files changed, 88 insertions(+), 2 deletions(-) create mode 100644 packages/backend/src/registries/ContainerRegistry.ts diff --git a/packages/backend/src/managers/playground.ts b/packages/backend/src/managers/playground.ts index 67e02d261..f10d9bf33 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 { 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.ts b/packages/backend/src/registries/ContainerRegistry.ts new file mode 100644 index 000000000..cd8e9aacc --- /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.type === '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.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); From b4cb16ae24ac632c4551c9f24ec6b42bf6536c3b Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Tue, 30 Jan 2024 15:11:38 +0100 Subject: [PATCH 2/7] test: adding unit tests Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- .../src/registries/ContainerRegistry.spec.ts | 123 ++++++++++++++++++ .../src/registries/ContainerRegistry.ts | 2 +- 2 files changed, 124 insertions(+), 1 deletion(-) create mode 100644 packages/backend/src/registries/ContainerRegistry.spec.ts diff --git a/packages/backend/src/registries/ContainerRegistry.spec.ts b/packages/backend/src/registries/ContainerRegistry.spec.ts new file mode 100644 index 000000000..9928ad4e6 --- /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 { 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 index cd8e9aacc..29aa7beb7 100644 --- a/packages/backend/src/registries/ContainerRegistry.ts +++ b/packages/backend/src/registries/ContainerRegistry.ts @@ -32,7 +32,7 @@ export class ContainerRegistry { 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.type === 'remove') { + if (event.status === 'remove') { this.subscribers.delete(event.id); } } From a762166ec92203725ac7cd93b80c790445893e90 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Tue, 30 Jan 2024 15:12:29 +0100 Subject: [PATCH 3/7] fix: linter Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- packages/backend/src/managers/playground.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/backend/src/managers/playground.ts b/packages/backend/src/managers/playground.ts index f10d9bf33..6c7de8f8e 100644 --- a/packages/backend/src/managers/playground.ts +++ b/packages/backend/src/managers/playground.ts @@ -32,7 +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 { ContainerRegistry } from '../registries/ContainerRegistry'; +import type { ContainerRegistry } from '../registries/ContainerRegistry'; // TODO: this should not be hardcoded const PLAYGROUND_IMAGE = 'quay.io/bootsy/playground:v0'; From 4e94c89fca22f09b753f9b0c19ba6b4662dbbad5 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Tue, 30 Jan 2024 15:22:10 +0100 Subject: [PATCH 4/7] fix: missing type in import Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- packages/backend/src/registries/ContainerRegistry.spec.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/backend/src/registries/ContainerRegistry.spec.ts b/packages/backend/src/registries/ContainerRegistry.spec.ts index 9928ad4e6..3e0eaac82 100644 --- a/packages/backend/src/registries/ContainerRegistry.spec.ts +++ b/packages/backend/src/registries/ContainerRegistry.spec.ts @@ -17,7 +17,7 @@ ***********************************************************************/ import { expect, test, vi } from 'vitest'; import { ContainerRegistry } from './ContainerRegistry'; -import { ContainerJSONEvent } from '@podman-desktop/api'; +import type { ContainerJSONEvent } from '@podman-desktop/api'; const mocks = vi.hoisted(() => ({ onEventMock: vi.fn(), From 2877fa86e027e53953a0675cf57eb38067becd6c Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Tue, 30 Jan 2024 16:11:44 +0100 Subject: [PATCH 5/7] test: fixing studio.spec.ts Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- packages/backend/src/studio.spec.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/backend/src/studio.spec.ts b/packages/backend/src/studio.spec.ts index f15bec5d8..d3624ef96 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(), + } }; }); From 0b63c0929ded280aa9dce6233a5b878a99b92b14 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Tue, 30 Jan 2024 16:13:36 +0100 Subject: [PATCH 6/7] fix: prettier Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- packages/backend/src/studio.spec.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/backend/src/studio.spec.ts b/packages/backend/src/studio.spec.ts index d3624ef96..bf7bb8e2a 100644 --- a/packages/backend/src/studio.spec.ts +++ b/packages/backend/src/studio.spec.ts @@ -48,7 +48,7 @@ vi.mock('@podman-desktop/api', async () => { }, containerEngine: { onEvent: vi.fn(), - } + }, }; }); From fd46312af6b041d5c09ac274062e739396630511 Mon Sep 17 00:00:00 2001 From: axel7083 <42176370+axel7083@users.noreply.github.com> Date: Tue, 30 Jan 2024 18:12:54 +0100 Subject: [PATCH 7/7] test: adapted to rebase Signed-off-by: axel7083 <42176370+axel7083@users.noreply.github.com> --- packages/backend/src/managers/playground.spec.ts | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) 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 () => {