feat: support screen capture for electron (tested in macOS, not teste…

…d in Windows)
Open-LLM-VTuber · Jan 28, 2025 · e543f70 · e543f70
1 parent c7e47db
commit e543f70
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 27 deletions.
diff --git a/src/main/index.ts b/src/main/index.ts
@@ -1,5 +1,5 @@
 /* eslint-disable no-shadow */
-import { app, ipcMain, globalShortcut } from "electron";
+import { app, ipcMain, globalShortcut, desktopCapturer } from "electron";
 import { electronApp, optimizer } from "@electron-toolkit/utils";
 import { WindowManager } from "./window-manager";
 import { MenuManager } from "./menu-manager";
@@ -56,6 +56,11 @@ function setupIPC(): void {
   ipcMain.on("update-config-files", (_event, files) => {
     menuManager.updateConfigFiles(files);
   });
+
+  ipcMain.handle('get-screen-capture', async () => {
+    const sources = await desktopCapturer.getSources({ types: ['screen'] });
+    return sources[0].id;
+  });
 }
 
 app.whenReady().then(() => {
@@ -81,18 +86,18 @@ app.whenReady().then(() => {
     return false;
   });
 
-  // if (process.env.NODE_ENV === "development") {
-  //   globalShortcut.register("F12", () => {
-  //     const window = windowManager.getWindow();
-  //     if (!window) return;
+  if (process.env.NODE_ENV === "development") {
+    globalShortcut.register("F12", () => {
+      const window = windowManager.getWindow();
+      if (!window) return;
 
-  //     if (window.webContents.isDevToolsOpened()) {
-  //       window.webContents.closeDevTools();
-  //     } else {
-  //       window.webContents.openDevTools();
-  //     }
-  //   });
-  // }
+      if (window.webContents.isDevToolsOpened()) {
+        window.webContents.closeDevTools();
+      } else {
+        window.webContents.openDevTools();
+      }
+    });
+  }
 
   setupIPC();
 
@@ -106,6 +111,16 @@ app.whenReady().then(() => {
   app.on("browser-window-created", (_, window) => {
     optimizer.watchWindowShortcuts(window);
   });
+
+  app.on('web-contents-created', (_, contents) => {
+    contents.session.setPermissionRequestHandler((webContents, permission, callback) => {
+      if (permission === 'media') {
+        callback(true);
+      } else {
+        callback(false);
+      }
+    });
+  });
 });
 
 app.on("window-all-closed", () => {

diff --git a/src/main/window-manager.ts b/src/main/window-manager.ts
@@ -45,7 +45,7 @@ export class WindowManager {
     });
   }
 
-  createWindow(p0: { titleBarOverlay: { color: string; symbolColor: string; height: number; }; }): BrowserWindow {
+  createWindow(options: Electron.BrowserWindowConstructorOptions): BrowserWindow {
     this.window = new BrowserWindow({
       width: 900,
       height: 670,
@@ -66,6 +66,7 @@ export class WindowManager {
       },
       hasShadow: false,
       paintWhenInitiallyHidden: true,
+      ...options,
     });
 
     this.setupWindowEvents();

diff --git a/src/preload/index.ts b/src/preload/index.ts
@@ -1,5 +1,5 @@
 /* eslint-disable @typescript-eslint/ban-ts-comment */
-import { contextBridge, ipcRenderer } from 'electron';
+import { contextBridge, ipcRenderer, desktopCapturer } from 'electron';
 import { electronAPI } from '@electron-toolkit/preload';
 import { ConfigFile } from '../main/menu-manager';
 
@@ -50,7 +50,20 @@ const api = {
 
 if (process.contextIsolated) {
   try {
-    contextBridge.exposeInMainWorld('electron', electronAPI);
+    contextBridge.exposeInMainWorld('electron', {
+      ...electronAPI,
+      desktopCapturer: {
+        getSources: (options) => desktopCapturer.getSources(options),
+      },
+      ipcRenderer: {
+        invoke: (channel, ...args) => ipcRenderer.invoke(channel, ...args),
+        on: (channel, func) => ipcRenderer.on(channel, func),
+        once: (channel, func) => ipcRenderer.once(channel, func),
+        removeListener: (channel, func) => ipcRenderer.removeListener(channel, func),
+        removeAllListeners: (channel) => ipcRenderer.removeAllListeners(channel),
+        send: (channel, ...args) => ipcRenderer.send(channel, ...args),
+      },
+    });
     contextBridge.exposeInMainWorld('api', api);
   } catch (error) {
     console.error(error);

diff --git a/src/renderer/src/context/screen-capture-context.tsx b/src/renderer/src/context/screen-capture-context.tsx
@@ -18,11 +18,36 @@ export function ScreenCaptureProvider({ children }: { children: ReactNode }) {
 
   const startCapture = async () => {
     try {
-      const displayMediaOptions: DisplayMediaStreamOptions = {
-        video: true,
-        audio: false,
-      };
-      const mediaStream = await navigator.mediaDevices.getDisplayMedia(displayMediaOptions);
+      let mediaStream: MediaStream;
+
+      if (window.electron) {
+        const sourceId = await window.electron.ipcRenderer.invoke('get-screen-capture');
+
+        const displayMediaOptions: DisplayMediaStreamOptions = {
+          video: {
+            // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+            // @ts-expect-error
+            mandatory: {
+              chromeMediaSource: "desktop",
+              chromeMediaSourceId: sourceId,
+              minWidth: 1280,
+              maxWidth: 1280,
+              minHeight: 720,
+              maxHeight: 720,
+            },
+          },
+          audio: false,
+        };
+
+        mediaStream = await navigator.mediaDevices.getUserMedia(displayMediaOptions);
+      } else {
+        const displayMediaOptions: DisplayMediaStreamOptions = {
+          video: true,
+          audio: false,
+        };
+        mediaStream = await navigator.mediaDevices.getDisplayMedia(displayMediaOptions);
+      }
+
       setStream(mediaStream);
       setIsStreaming(true);
       setError('');

diff --git a/src/renderer/src/hooks/utils/use-media-capture.tsx b/src/renderer/src/hooks/utils/use-media-capture.tsx
@@ -1,6 +1,7 @@
 import { useCallback } from 'react';
 import { useCamera } from '@/context/camera-context';
 import { useScreenCaptureContext } from '@/context/screen-capture-context';
+import { toaster } from "@/components/ui/toaster";
 
 // Add type definition for ImageCapture
 declare class ImageCapture {
@@ -19,11 +20,17 @@ export function useMediaCapture() {
   const { stream: cameraStream } = useCamera();
   const { stream: screenStream } = useScreenCaptureContext();
 
-  const captureFrame = useCallback(async (stream: MediaStream | null) => {
-    if (!stream) return null;
+  const captureFrame = useCallback(async (stream: MediaStream | null, source: 'camera' | 'screen') => {
+    if (!stream) {
+      console.warn(`No ${source} stream available`);
+      return null;
+    }
 
     const videoTrack = stream.getVideoTracks()[0];
-    if (!videoTrack) return null;
+    if (!videoTrack) {
+      console.warn(`No video track in ${source} stream`);
+      return null;
+    }
 
     const imageCapture = new ImageCapture(videoTrack);
     try {
@@ -32,12 +39,20 @@ export function useMediaCapture() {
       canvas.width = bitmap.width;
       canvas.height = bitmap.height;
       const ctx = canvas.getContext('2d');
-      if (!ctx) return null;
+      if (!ctx) {
+        console.error('Failed to get canvas context');
+        return null;
+      }
 
       ctx.drawImage(bitmap, 0, 0);
       return canvas.toDataURL('image/jpeg', 0.8);
     } catch (error) {
-      console.error('Error capturing frame:', error);
+      console.error(`Error capturing ${source} frame:`, error);
+      toaster.create({
+        title: `Failed to capture ${source} frame: ${error}`,
+        type: 'error',
+        duration: 2000,
+      });
       return null;
     }
   }, []);
@@ -47,7 +62,7 @@ export function useMediaCapture() {
 
     // Capture camera frame
     if (cameraStream) {
-      const cameraFrame = await captureFrame(cameraStream);
+      const cameraFrame = await captureFrame(cameraStream, 'camera');
       if (cameraFrame) {
         images.push({
           source: 'camera',
@@ -59,7 +74,7 @@ export function useMediaCapture() {
 
     // Capture screen frame
     if (screenStream) {
-      const screenFrame = await captureFrame(screenStream);
+      const screenFrame = await captureFrame(screenStream, 'screen');
       if (screenFrame) {
         images.push({
           source: 'screen',