fix: ensure screen capture uses primary display source (close: #116)

src/main/agent/device.test.ts

@@ -0,0 +1,76 @@
+import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest';
+import { screen, desktopCapturer } from 'electron';
+import { Desktop } from './device';
+
+// Mock dependencies
+vi.mock('electron', () => ({
+  screen: {
+    getPrimaryDisplay: vi.fn(),
+  },
+  desktopCapturer: {
+    getSources: vi.fn(),
+  },
+  app: {
+    on: vi.fn(),
+    off: vi.fn(),
+    quit: vi.fn(),
+    exit: vi.fn(),
+    relaunch: vi.fn(),
+  },
+}));
+
+vi.mock('./execute', () => ({
+  execute: vi.fn(),
+}));
+
+vi.mock('@main/env', () => ({
+  isMacOS: false,
+}));
+
+describe('Desktop', () => {
+  let desktop: Desktop;
+
+  beforeEach(() => {
+    desktop = new Desktop();
+    vi.clearAllMocks();
+  });
+
+  afterEach(() => {
+    vi.resetAllMocks();
+  });
+
+  describe('screenshot', () => {
+    it('should capture screenshot successfully', async () => {
+      const mockDisplay = {
+        id: '1',
+        size: { width: 1920, height: 1080 },
+      };
+      const mockSource = {
+        display_id: '1',
+        thumbnail: {
+          toPNG: () => Buffer.from('mock-image'),
+        },
+      };
+
+      vi.mocked(screen.getPrimaryDisplay).mockReturnValue(mockDisplay as any);
+      vi.mocked(desktopCapturer.getSources).mockResolvedValueOnce([
+        mockSource as any,
+      ]);
+
+      const result = await desktop.screenshot();
+
+      expect(result).toEqual({
+        base64: 'bW9jay1pbWFnZQ==',
+        width: 1920,
+        height: 1080,
+      });
+      expect(desktopCapturer.getSources).toHaveBeenCalledWith({
+        types: ['screen'],
+        thumbnailSize: {
+          width: 1920,
+          height: 1080,
+        },
+      });
+    });
+  });
+});

src/main/agent/device.ts

@@ -81,8 +81,10 @@ export class Desktop {
         height: Math.round(height),
       },
     });
-    const primarySource = sources[0];
-    const screenshot = primarySource.thumbnail;
+    const primarySource = sources.find(
+      (source) => source.display_id === primaryDisplay.id.toString(),
+    );
+    const screenshot = primarySource!.thumbnail;
 
     return {
       base64: screenshot.toPNG().toString('base64'),

src/main/main.ts

@@ -12,6 +12,7 @@
   ipcMain,
   session,
   WebContentsView,
+  screen,
 } from 'electron';
 import squirrelStartup from 'electron-squirrel-startup';
 import ElectronStore from 'electron-store';
@@ -131,8 +132,12 @@
   session.defaultSession.setDisplayMediaRequestHandler(
     (_request, callback) => {
       desktopCapturer.getSources({ types: ['screen'] }).then((sources) => {
-        // Grant access to the first screen found.
-        callback({ video: sources[0], audio: 'loopback' });
+        const primaryDisplay = screen.getPrimaryDisplay();
+        const primarySource = sources.find(
+          (source) => source.display_id === primaryDisplay.id.toString(),
+        );
+
+        callback({ video: primarySource!, audio: 'loopback' });
       });
     },
     { useSystemPicker: false },