feat: listen to podman desktop container api events

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 () => {

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<string, PlaygroundState>;
   private queries: Map<number, QueryState>;
 
-  constructor(private webview: Webview) {
+  constructor(
+    private webview: Webview,
+    private containerRegistry: ContainerRegistry,
+  ) {
     this.playgrounds = new Map<string, PlaygroundState>();
     this.queries = new Map<number, QueryState>();
   }
@@ -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,

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);
+});

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<string, Subscriber[]> = 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),
+      );
+    });
+  }
+}

packages/backend/src/studio.spec.ts

@@ -46,6 +46,9 @@ vi.mock('@podman-desktop/api', async () => {
         },
       }),
     },
+    containerEngine: {
+      onEvent: vi.fn(),
+    },
   };
 });
 

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,14 +95,18 @@ 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);
 
     this.rpcExtension = new RpcExtension(this.#panel.webview);
     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);