feat: listen podman machine stop to mark playgrounds as stopped (#228)

* feat: listen machine stop to detect playground stopped

* test startupSubscribe

* unit tests for listenMachine

* unit tests for adoptRunningPlaygrounds

packages/backend/src/managers/playground.spec.ts

@@ -17,10 +17,10 @@
  ***********************************************************************/
 
 import { beforeEach, afterEach, expect, test, vi } from 'vitest';
-import { PlayGroundManager } from './playground';
+import { LABEL_MODEL_ID, LABEL_MODEL_PORT, PlayGroundManager } from './playground';
 import type { ImageInfo, Webview } from '@podman-desktop/api';
 import type { ContainerRegistry } from '../registries/ContainerRegistry';
-import type { PodmanConnection } from './podmanConnection';
+import type { PodmanConnection, machineStopHandle, startupHandle } from './podmanConnection';
 
 const mocks = vi.hoisted(() => ({
   postMessage: vi.fn(),
@@ -30,6 +30,9 @@ const mocks = vi.hoisted(() => ({
   stopContainer: vi.fn(),
   getFreePort: vi.fn(),
   containerRegistrySubscribeMock: vi.fn(),
+  startupSubscribe: vi.fn(),
+  onMachineStop: vi.fn(),
+  listContainers: vi.fn(),
 }));
 
 vi.mock('@podman-desktop/api', async () => {
@@ -41,6 +44,7 @@ vi.mock('@podman-desktop/api', async () => {
       pullImage: mocks.pullImage,
       createContainer: mocks.createContainer,
       stopContainer: mocks.stopContainer,
+      listContainers: mocks.listContainers,
     },
   };
 });
@@ -66,7 +70,10 @@ beforeEach(() => {
       postMessage: mocks.postMessage,
     } as unknown as Webview,
     containerRegistryMock,
-    {} as PodmanConnection,
+    {
+      startupSubscribe: mocks.startupSubscribe,
+      onMachineStop: mocks.onMachineStop,
+    } as unknown as PodmanConnection,
   );
   originalFetch = globalThis.fetch;
   globalThis.fetch = vi.fn().mockResolvedValue({});
@@ -165,3 +172,42 @@ test('stopPlayground should stop a started playground', async () => {
   await manager.stopPlayground('model1');
   expect(mocks.stopContainer).toHaveBeenNthCalledWith(1, 'engine1', 'container1');
 });
+
+test('adoptRunningPlaygrounds updates the playground state with the found container', async () => {
+  mocks.listContainers.mockResolvedValue([
+    {
+      Id: 'container-id-1',
+      engineId: 'engine-id-1',
+      Labels: {
+        [LABEL_MODEL_ID]: 'model-id-1',
+        [LABEL_MODEL_PORT]: '8080',
+      },
+      State: 'running',
+    },
+  ]);
+  mocks.startupSubscribe.mockImplementation((f: startupHandle) => {
+    f();
+  });
+  const updatePlaygroundStateSpy = vi.spyOn(manager, 'updatePlaygroundState');
+  manager.adoptRunningPlaygrounds();
+  await new Promise(resolve => setTimeout(resolve, 0));
+  expect(updatePlaygroundStateSpy).toHaveBeenNthCalledWith(1, 'model-id-1', {
+    container: {
+      containerId: 'container-id-1',
+      engineId: 'engine-id-1',
+      port: 8080,
+    },
+    modelId: 'model-id-1',
+    status: 'running',
+  });
+});
+
+test('onMachineStop updates the playground state with no playground running', async () => {
+  mocks.listContainers.mockResolvedValue([]);
+  mocks.onMachineStop.mockImplementation((f: machineStopHandle) => {
+    f();
+  });
+  const sendPlaygroundStateSpy = vi.spyOn(manager, 'sendPlaygroundState').mockResolvedValue();
+  manager.adoptRunningPlaygrounds();
+  expect(sendPlaygroundStateSpy).toHaveBeenCalledOnce();
+});

packages/backend/src/managers/playground.ts

@@ -35,8 +35,8 @@ import type { PodmanConnection } from './podmanConnection';
 import OpenAI from 'openai';
 import { timeout } from '../utils/utils';
 
-const LABEL_MODEL_ID = 'ai-studio-model-id';
-const LABEL_MODEL_PORT = 'ai-studio-model-port';
+export const LABEL_MODEL_ID = 'ai-studio-model-id';
+export const LABEL_MODEL_PORT = 'ai-studio-model-port';
 
 // TODO: this should not be hardcoded
 const PLAYGROUND_IMAGE = 'quay.io/bootsy/playground:v0';
@@ -100,6 +100,11 @@ export class PlayGroundManager {
           console.error('error during adoption of existing playground containers', err);
         });
     });
+    this.podmanConnection.onMachineStop(() => {
+      // Podman Machine has been stopped, we consider all playground containers are stopped
+      this.playgrounds.clear();
+      this.sendPlaygroundState();
+    });
   }
 
   async selectImage(image: string): Promise<ImageInfo | undefined> {
@@ -117,6 +122,10 @@ export class PlayGroundManager {
 
   updatePlaygroundState(modelId: string, state: PlaygroundState): void {
     this.playgrounds.set(modelId, state);
+    this.sendPlaygroundState();
+  }
+
+  sendPlaygroundState() {
     this.webview
       .postMessage({
         id: MSG_PLAYGROUNDS_STATE_UPDATE,

packages/backend/src/managers/podmanConnection.spec.ts

@@ -0,0 +1,113 @@
+import { expect, test, vi } from 'vitest';
+import { PodmanConnection } from './podmanConnection';
+import type { RegisterContainerConnectionEvent, UpdateContainerConnectionEvent } from '@podman-desktop/api';
+
+const mocks = vi.hoisted(() => ({
+  getContainerConnections: vi.fn(),
+  onDidRegisterContainerConnection: vi.fn(),
+  onDidUpdateContainerConnection: vi.fn(),
+}));
+
+vi.mock('@podman-desktop/api', async () => {
+  return {
+    provider: {
+      onDidRegisterContainerConnection: mocks.onDidRegisterContainerConnection,
+      getContainerConnections: mocks.getContainerConnections,
+      onDidUpdateContainerConnection: mocks.onDidUpdateContainerConnection,
+    },
+  };
+});
+
+test('startupSubscribe should execute immediately if provider already registered', () => {
+  const manager = new PodmanConnection();
+  // one provider is already registered
+  mocks.getContainerConnections.mockReturnValue([
+    {
+      connection: {
+        type: 'podman',
+        status: () => 'started',
+      },
+    },
+  ]);
+  mocks.onDidRegisterContainerConnection.mockReturnValue({
+    dispose: vi.fn,
+  });
+  manager.listenRegistration();
+  const handler = vi.fn();
+  manager.startupSubscribe(handler);
+  // the handler is called immediately
+  expect(handler).toHaveBeenCalledOnce();
+});
+
+test('startupSubscribe should execute  when provider is registered', async () => {
+  const manager = new PodmanConnection();
+
+  // no provider is already registered
+  mocks.getContainerConnections.mockReturnValue([]);
+  mocks.onDidRegisterContainerConnection.mockImplementation((f: (e: RegisterContainerConnectionEvent) => void) => {
+    setTimeout(() => {
+      f({
+        connection: {
+          type: 'podman',
+          status: () => 'started',
+        },
+      } as unknown as RegisterContainerConnectionEvent);
+    }, 1);
+    return {
+      dispose: vi.fn(),
+    };
+  });
+  manager.listenRegistration();
+  const handler = vi.fn();
+  manager.startupSubscribe(handler);
+  // the handler is not called immediately
+  expect(handler).not.toHaveBeenCalledOnce();
+  await new Promise(resolve => setTimeout(resolve, 10));
+  expect(handler).toHaveBeenCalledOnce();
+});
+
+test('onMachineStart should call the handler when machine starts', async () => {
+  const manager = new PodmanConnection();
+  mocks.onDidUpdateContainerConnection.mockImplementation((f: (e: UpdateContainerConnectionEvent) => void) => {
+    setTimeout(() => {
+      f({
+        connection: {
+          type: 'podman',
+        },
+        status: 'started',
+      } as UpdateContainerConnectionEvent);
+    }, 1);
+    return {
+      dispose: vi.fn(),
+    };
+  });
+  manager.listenMachine();
+  const handler = vi.fn();
+  manager.onMachineStart(handler);
+  expect(handler).not.toHaveBeenCalledOnce();
+  await new Promise(resolve => setTimeout(resolve, 10));
+  expect(handler).toHaveBeenCalledOnce();
+});
+
+test('onMachineStop should call the handler when machine stops', async () => {
+  const manager = new PodmanConnection();
+  mocks.onDidUpdateContainerConnection.mockImplementation((f: (e: UpdateContainerConnectionEvent) => void) => {
+    setTimeout(() => {
+      f({
+        connection: {
+          type: 'podman',
+        },
+        status: 'stopped',
+      } as UpdateContainerConnectionEvent);
+    }, 1);
+    return {
+      dispose: vi.fn(),
+    };
+  });
+  manager.listenMachine();
+  const handler = vi.fn();
+  manager.onMachineStop(handler);
+  expect(handler).not.toHaveBeenCalledOnce();
+  await new Promise(resolve => setTimeout(resolve, 10));
+  expect(handler).toHaveBeenCalledOnce();
+});

packages/backend/src/managers/podmanConnection.ts

@@ -16,15 +16,28 @@
  * SPDX-License-Identifier: Apache-2.0
  ***********************************************************************/
 
-import { type RegisterContainerConnectionEvent, provider } from '@podman-desktop/api';
+import {
+  type RegisterContainerConnectionEvent,
+  provider,
+  type UpdateContainerConnectionEvent,
+} from '@podman-desktop/api';
 
-type startupHandle = () => void;
+export type startupHandle = () => void;
+export type machineStartHandle = () => void;
+export type machineStopHandle = () => void;
 
 export class PodmanConnection {
   #firstFound = false;
   #toExecuteAtStartup: startupHandle[] = [];
+  #toExecuteAtMachineStop: machineStopHandle[] = [];
+  #toExecuteAtMachineStart: machineStartHandle[] = [];
 
   init(): void {
+    this.listenRegistration();
+    this.listenMachine();
+  }
+
+  listenRegistration() {
     // In case the extension has not yet registered, we listen for new registrations
     // and retain the first started podman provider
     const disposable = provider.onDidRegisterContainerConnection((e: RegisterContainerConnectionEvent) => {
@@ -62,4 +75,29 @@ export class PodmanConnection {
       this.#toExecuteAtStartup.push(f);
     }
   }
+
+  listenMachine() {
+    provider.onDidUpdateContainerConnection((e: UpdateContainerConnectionEvent) => {
+      if (e.connection.type !== 'podman') {
+        return;
+      }
+      if (e.status === 'stopped') {
+        for (const f of this.#toExecuteAtMachineStop) {
+          f();
+        }
+      } else if (e.status === 'started') {
+        for (const f of this.#toExecuteAtMachineStart) {
+          f();
+        }
+      }
+    });
+  }
+
+  onMachineStart(f: machineStartHandle) {
+    this.#toExecuteAtMachineStart.push(f);
+  }
+
+  onMachineStop(f: machineStopHandle) {
+    this.#toExecuteAtMachineStop.push(f);
+  }
 }

packages/backend/src/studio.spec.ts

@@ -57,6 +57,7 @@ vi.mock('@podman-desktop/api', async () => {
     },
     provider: {
       onDidRegisterContainerConnection: vi.fn(),
+      onDidUpdateContainerConnection: vi.fn(),
       getContainerConnections: mocks.getContainerConnections,
     },
   };