add support for kde(wayland sessions)

memento/background.py

@@ -6,7 +6,9 @@
 import memento.utils as utils
 import asyncio
 import os
+import pyscreenshot as ImageGrab
 import time
+import threading
 import multiprocessing
 from multiprocessing import Queue
 import signal
@@ -16,7 +18,7 @@
 from memento.db import Db
 from langchain.vectorstores import Chroma
 from memento.segments import AppSegments
-
+import memento.kwin as kwin
 
 class Background:
     def __init__(self):
@@ -58,8 +60,22 @@ def __init__(self):
                 embedding_function=OpenAIEmbeddings(),
                 collection_name="memento_db",
             )
+        # KDE or X11 tracking
+        xdg_session_desktop = os.environ.get("XDG_SESSION_DESKTOP", "").lower()
+        xdg_current_desktop = os.environ.get("XDG_CURRENT_DESKTOP", "").lower()
+        xdg_session_type = os.environ.get("XDG_SESSION_TYPE", "").lower()
+
+        is_kde = "kde" in xdg_session_desktop or "kde" in xdg_current_desktop
+        is_wayland = xdg_session_type == "wayland"
 
-        self.sct = mss.mss()
+        if is_kde and is_wayland:
+            self.stop_event = threading.Event()
+            self.watcher_thread = threading.Thread(target=self.start_watcher)
+            self.watcher_thread.start()
+            self.method = "kde"
+        else:
+            self.method = "x11"
+            self.sct = mss.mss()
         self.rec = utils.Recorder(
             os.path.join(utils.CACHE_PATH, str(self.nb_rec) + ".mp4")
         )
@@ -82,6 +98,10 @@ def __init__(self):
             self.workers[i].start()
             print("started worker", i)
 
+    def start_watcher(self):
+        watcher = kwin.WindowWatcher(self.stop_event)
+        asyncio.run(watcher.run())
+
     def process_images(self):
         # Infinite worker
 
@@ -121,6 +141,10 @@ def process_images(self):
     def stop_rec(self, sig, frame):
         # self.rec.stop()
         print("STOPPING MAIN", os.getpid())
+        if self.method == "kde":
+            self.stop_event.set()  # Signal the watcher to stop
+            self.watcher_thread.join()  # Wait for the watcher to finish
+            print("🛑 KWin Window Watcher stopped.")
         exit()
 
     def stop_process(self, sig, frame):
@@ -135,12 +159,18 @@ def run(self):
             (utils.RESOLUTION[1], utils.RESOLUTION[0], 3), dtype=np.uint8
         )
         while self.running:
-            window_title = utils.get_active_window()
-
-            # Get screenshot and add it to recorder
-            im = np.array(self.sct.grab(self.sct.monitors[1]))
-            im = im[:, :, :-1]
+            if self.method == "kde" :
+                window_title = kwin.active_window.resource_class
+                im = np.array(ImageGrab.grab())
+                # Reverse the order of the color channels
+                # Swap RGB to BGR
+                im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
+            else: 
+                window_title = utils.get_active_window()
+                im = np.array(self.sct.grab(self.sct.monitors[1]))
+                im = im[:, :, :-1]
             im = cv2.resize(im, utils.RESOLUTION)
+            # Get screenshot and add it to recorder
             asyncio.run(self.rec.new_im(im))
 
             # Create metadata

memento/kwin.py

@@ -0,0 +1,193 @@
+from dbus_fast.aio import MessageBus
+from dbus_fast.service import ServiceInterface, method
+import os
+import asyncio
+from pathlib import Path
+
+# Constants
+BUS_NAME = "com.apirrone.Memento"
+OBJECT_PATH = "/com/apirrone/Memento"
+KWIN_SCRIPT_PATH = Path("/tmp/kwin_window.js")
+
+# Active Window class (Stores the active window details)
+class ActiveWindow:
+    def __init__(self):
+        self.caption = ""
+        self.resource_class = ""
+        self.resource_name = ""
+
+active_window = ActiveWindow()
+
+# D-Bus Interface for Active Window Notifications
+class ActiveWindowInterface(ServiceInterface):
+    def __init__(self):
+        super().__init__(BUS_NAME)
+
+    @method()
+    def NotifyActiveWindow(self, caption: 's', resource_class: 's', resource_name: 's') -> 's':
+        print(f"\n📌 Active Window Updated:\nCaption: {caption}\nClass: {resource_class}\nName: {resource_name}")
+        active_window.caption = caption
+        active_window.resource_class = resource_class
+        active_window.resource_name = resource_name
+
+# KWin Script Manager (Handles loading, unloading, and starting KWin scripts)
+
+class KWinScript:
+
+    async def load(self):
+        print("📂 Loading KWin script...")
+
+        kwin_script_code = """
+        let connections = {};
+        function send(client) {
+            callDBus(
+                "com.apirrone.Memento",
+                "/com/apirrone/Memento",
+                "com.apirrone.Memento",
+                "NotifyActiveWindow",
+                "caption" in client ? client.caption : "",
+                "resourceClass" in client ? String(client.resourceClass) : "",
+                "resourceName" in client ? String(client.resourceName) : ""
+            );
+        }
+        let handler = function(client){
+            if (client === null) {
+                return;
+            }
+            if (!(client.internalId in connections)) {
+                connections[client.internalId] = true;
+                client.captionChanged.connect(function() {
+                    if (client.active) {
+                        send(client);
+                    }
+                });
+            }
+
+            send(client);
+        };
+
+        let activationEvent = workspace.windowActivated ? workspace.windowActivated : workspace.clientActivated;
+        if (workspace.windowActivated) {
+            workspace.windowActivated.connect(handler);
+        } else {
+            // KDE version < 6
+            workspace.clientActivated.connect(handler);
+        }
+        """
+        KWIN_SCRIPT_PATH.write_text(kwin_script_code)
+
+        bus = await MessageBus().connect()
+
+        # Retrieve introspection data
+        introspection_data = await bus.introspect("org.kde.KWin", "/Scripting")
+
+        # Obtain a proxy object
+        proxy_object = bus.get_proxy_object("org.kde.KWin", "/Scripting", introspection_data)
+
+        # Get the interface
+        kwin_interface = proxy_object.get_interface("org.kde.kwin.Scripting")
+
+        script_number = await kwin_interface.call_load_script(str(KWIN_SCRIPT_PATH))
+        await self.start(script_number)
+        KWIN_SCRIPT_PATH.unlink()
+
+    async def is_script_loaded(self):
+        bus = await MessageBus().connect()
+
+        # Retrieve introspection data
+        introspection_data = await bus.introspect("org.kde.KWin", "/Scripting")
+
+        # Obtain a proxy object
+        proxy_object = bus.get_proxy_object("org.kde.KWin", "/Scripting", introspection_data)
+
+        # Get the interface
+        kwin_interface = proxy_object.get_interface("org.kde.kwin.Scripting")
+
+        # Call the isScriptLoaded method
+        result = await kwin_interface.call_is_script_loaded(str(KWIN_SCRIPT_PATH))
+        return result
+
+    async def start(self, script_number):
+        print(f"🚀 Starting KWin script {script_number}...")
+        kwin_path = f"/Scripting/Script{script_number}" if await self.get_major_version() >= 6 else f"/{script_number}"
+        bus = await MessageBus().connect()
+
+        # Retrieve introspection data
+        introspection_data = await bus.introspect("org.kde.KWin", kwin_path)
+
+        # Obtain a proxy object
+        proxy_object = bus.get_proxy_object("org.kde.KWin", kwin_path, introspection_data)
+
+        # Get the interface
+        kwin_interface = proxy_object.get_interface("org.kde.kwin.Script")
+        await kwin_interface.call_run()
+
+    async def unload(self):
+        print("🛑 Unloading KWin script...")
+        bus = await MessageBus().connect()
+
+        # Retrieve introspection data
+        introspection_data = await bus.introspect("org.kde.KWin", "/Scripting")
+
+        # Obtain a proxy object
+        proxy_object = bus.get_proxy_object("org.kde.KWin", "/Scripting", introspection_data)
+
+        # Get the interface
+        kwin_interface = proxy_object.get_interface("org.kde.kwin.Scripting")
+
+        # Call the unloadScript method
+        await kwin_interface.call_unload_script(str(KWIN_SCRIPT_PATH))
+        print("✅ KWin script successfully unloaded.")
+
+    async def get_major_version(self):
+        kde_version = os.getenv("KDE_SESSION_VERSION")
+        if kde_version:
+            return int(kde_version)
+
+        try:
+            bus = await MessageBus().connect()
+
+            # Retrieve introspection data
+            introspection_data = await bus.introspect("org.kde.KWin", "/KWin")
+
+            # Obtain a proxy object
+            proxy_object = bus.get_proxy_object("org.kde.KWin", "/KWin", introspection_data)
+
+            # Get the interface
+            kwin_interface = proxy_object.get_interface("org.kde.KWin")
+            support_info = await kwin_interface.call_support_information()
+
+            for line in support_info.splitlines():
+                if "KWin version:" in line:
+                    return int(line.split()[2].split(".")[0])  # Extract major version
+        except Exception as e:
+            print(f"⚠️ Failed to retrieve KDE version from D-Bus: {e}")
+
+        return 5  # Default to KDE 5 if not found
+
+# Window Watcher Class (Manages Active Window tracking and D-Bus interface)
+class WindowWatcher:
+    def __init__(self, stop_event):
+        self.kwin_script = KWinScript()
+        self.stop_event = stop_event
+
+    async def run(self):
+        if await self.kwin_script.is_script_loaded():
+            print("⚠️ KWin script already loaded. Unloading and reloading...")
+            await self.kwin_script.unload()
+
+        bus = await MessageBus().connect()
+        interface = ActiveWindowInterface()
+        bus.export(OBJECT_PATH, interface)
+        await bus.request_name(BUS_NAME)
+
+        print(f"✅ D-Bus Service Running: {BUS_NAME} on {OBJECT_PATH}")
+        await self.kwin_script.load()
+
+        try:
+            while not self.stop_event.is_set():
+                await asyncio.sleep(1)  # Sleep briefly to yield control
+        except asyncio.CancelledError:
+            print("\nStopping Window Watcher...")
+            await self.kwin_script.unload()
+

setup.py

@@ -16,6 +16,8 @@
         "opencv-contrib-python==4.8.0.74",
         "xlib==0.21",
         "av==10.0.0",
+        "dbus-fast==2.33.0",
+        "pyscreenshot==3.1",
         "pygame==2.5.0",
         "TextTron==0.45",
         "thefuzz==0.19.0",