windows-mcp 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

windows_mcp/desktop/service.py

@@ -1,5 +1,6 @@
-from windows_mcp.desktop.config import BROWSER_NAMES, PROCESS_PER_MONITOR_DPI_AWARE
-from windows_mcp.desktop.views import DesktopState, App, Status, Size
+from windows_mcp.vdm.core import get_all_desktops, get_current_desktop, is_window_on_current_desktop
+from windows_mcp.desktop.views import DesktopState, Window, Browser, Status, Size
+from windows_mcp.desktop.config import PROCESS_PER_MONITOR_DPI_AWARE
 from windows_mcp.tree.views import BoundingBox, TreeElementNode
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from PIL import ImageGrab, ImageFont, ImageDraw, Image
@@ -45,25 +46,36 @@ class Desktop:
         self.tree=Tree(self)
         self.desktop_state=None
         
-    def get_state(self,use_annotation:bool=True,use_vision:bool=False,use_dom:bool=False,as_bytes:bool=False,scale:float=1.0)->DesktopState:
-        sleep(0.1)
+    def get_state(self,use_annotation:bool|str=True,use_vision:bool|str=False,use_dom:bool|str=False,as_bytes:bool|str=False,scale:float=1.0)->DesktopState:
+        use_annotation = use_annotation is True or (isinstance(use_annotation, str) and use_annotation.lower() == 'true')
+        use_vision = use_vision is True or (isinstance(use_vision, str) and use_vision.lower() == 'true')
+        use_dom = use_dom is True or (isinstance(use_dom, str) and use_dom.lower() == 'true')
+        as_bytes = as_bytes is True or (isinstance(as_bytes, str) and as_bytes.lower() == 'true')
+        
         start_time = time()
 
         controls_handles=self.get_controls_handles() # Taskbar,Program Manager,Apps, Dialogs
-        apps,apps_handles=self.get_apps(controls_handles=controls_handles) # Apps
-        active_app=self.get_active_app(apps=apps) #Active App
-        active_app_handle=active_app.handle if active_app else None
+        windows,windows_handles=self.get_windows(controls_handles=controls_handles) # Apps
+        active_window=self.get_active_window(windows=windows) #Active Window
+        active_window_handle=active_window.handle if active_window else None
+
+        try:
+            active_desktop=get_current_desktop()
+            all_desktops=get_all_desktops()
+        except RuntimeError:
+            active_desktop = {'id': '00000000-0000-0000-0000-000000000000', 'name': 'Default Desktop'}
+            all_desktops = [active_desktop]
 
-        if active_app is not None and active_app in apps:
-            apps.remove(active_app)
+        if active_window is not None and active_window in windows:
+            windows.remove(active_window)
 
-        logger.debug(f"Active app: {active_app or 'No Active App Found'}")
-        logger.debug(f"Apps: {apps}")
+        logger.debug(f"Active window: {active_window or 'No Active Window Found'}")
+        logger.debug(f"Windows: {windows}")
         
         #Preparing handles for Tree
-        other_apps_handles=list(controls_handles-apps_handles)
+        other_windows_handles=list(controls_handles-windows_handles)
 
-        tree_state=self.tree.get_state(active_app_handle,other_apps_handles,use_dom=use_dom)
+        tree_state=self.tree.get_state(active_window_handle,other_windows_handles,use_dom=use_dom)
 
         if use_vision:
             if use_annotation:
@@ -83,13 +95,20 @@ class Desktop:
         else:
             screenshot=None
             
-        self.desktop_state=DesktopState(apps= apps,active_app=active_app,screenshot=screenshot,tree_state=tree_state)
+        self.desktop_state=DesktopState(
+            active_window=active_window,
+            windows=windows,
+            active_desktop=active_desktop,
+            all_desktops=all_desktops,
+            screenshot=screenshot,
+            tree_state=tree_state
+        )
         # Log the time taken to capture the state
         end_time = time()
         logger.info(f"Desktop State capture took {end_time - start_time:.2f} seconds")
         return self.desktop_state
     
-    def get_app_status(self,control:uia.Control)->Status:
+    def get_window_status(self,control:uia.Control)->Status:
         if uia.IsIconic(control.NativeWindowHandle):
             return Status.MINIMIZED
         elif uia.IsZoomed(control.NativeWindowHandle):
@@ -117,8 +136,8 @@ class Desktop:
         try:
             reader = csv.DictReader(io.StringIO(apps_info.strip()))
             return {
-                row.get('Name').lower(): row.get('AppID') 
-                for row in reader 
+                row.get('Name', '').lower(): row.get('AppID', '')
+                for row in reader
                 if row.get('Name') and row.get('AppID')
             }
         except Exception as e:
@@ -129,10 +148,11 @@ class Desktop:
         try:
             encoded = base64.b64encode(command.encode("utf-16le")).decode("ascii")
             result = subprocess.run(
-                ['powershell', '-NoProfile', '-EncodedCommand', encoded], 
+                ['powershell', '-NoProfile', '-OutputFormat', 'Text', '-EncodedCommand', encoded], 
                 capture_output=True,  # No errors='ignore' - let subprocess return bytes
                 timeout=timeout,
-                cwd=os.path.expanduser(path='~')
+                cwd=os.path.expanduser(path='~'),
+                env=os.environ.copy()  # Inherit environment variables including PATH
             )
             # Handle both bytes and str output (subprocess behavior varies by environment)
             stdout = result.stdout
@@ -147,11 +167,11 @@ class Desktop:
         except Exception as e:
             return (f'Command execution failed: {type(e).__name__}: {e}', 1)
         
-    def is_app_browser(self,node:uia.Control):
+    def is_window_browser(self,node:uia.Control):
         '''Give any node of the app and it will return True if the app is a browser, False otherwise.'''
         try:
             process=Process(node.ProcessId)
-            return process.name() in BROWSER_NAMES
+            return Browser.has_process(process.name())
         except:
             return False
     
@@ -162,32 +182,32 @@ class Desktop:
         return "".join([row.get('DisplayName') for row in reader])
     
     def resize_app(self,size:tuple[int,int]=None,loc:tuple[int,int]=None)->tuple[str,int]:
-        active_app=self.desktop_state.active_app
-        if active_app is None:
-            return "No active app found",1
-        if active_app.status==Status.MINIMIZED:
-            return f"{active_app.name} is minimized",1
-        elif active_app.status==Status.MAXIMIZED:
-            return f"{active_app.name} is maximized",1
+        active_window=self.desktop_state.active_window
+        if active_window is None:
+            return "No active window found",1
+        if active_window.status==Status.MINIMIZED:
+            return f"{active_window.name} is minimized",1
+        elif active_window.status==Status.MAXIMIZED:
+            return f"{active_window.name} is maximized",1
         else:
-            app_control=uia.ControlFromHandle(active_app.handle)
+            window_control=uia.ControlFromHandle(active_window.handle)
             if loc is None:
-                x=app_control.BoundingRectangle.left
-                y=app_control.BoundingRectangle.top
+                x=window_control.BoundingRectangle.left
+                y=window_control.BoundingRectangle.top
                 loc=(x,y)
             if size is None:
-                width=app_control.BoundingRectangle.width()
-                height=app_control.BoundingRectangle.height()
+                width=window_control.BoundingRectangle.width()
+                height=window_control.BoundingRectangle.height()
                 size=(width,height)
             x,y=loc
             width,height=size
-            app_control.MoveWindow(x,y,width,height)
-            return (f'{active_app.name} resized to {width}x{height} at {x},{y}.',0)
+            window_control.MoveWindow(x,y,width,height)
+            return (f'{active_window.name} resized to {width}x{height} at {x},{y}.',0)
     
     def is_app_running(self,name:str)->bool:
-        apps, _ = self.get_apps()
-        apps_dict = {app.name: app for app in apps}
-        return process.extractOne(name,list(apps_dict.keys()),score_cutoff=60) is not None
+        windows, _ = self.get_windows()
+        windows_dict = {window.name: window for window in windows}
+        return process.extractOne(name,list(windows_dict.keys()),score_cutoff=60) is not None
     
     def app(self,mode:Literal['launch','switch','resize'],name:Optional[str]=None,loc:Optional[tuple[int,int]]=None,size:Optional[tuple[int,int]]=None):
         match mode:
@@ -207,7 +227,7 @@ class Desktop:
                     safe_name = re.escape(name)
                     if uia.WindowControl(RegexName=f'(?i).*{safe_name}.*').Exists(maxSearchSeconds=10):
                         launched = True
-
+ 
                 if launched:
                     return f'{name.title()} launched.'
                 return f'Launching {name.title()} sent, but window not detected yet.'
@@ -232,37 +252,41 @@ class Desktop:
         app_name,_=matched_app
         appid=apps_map.get(app_name)
         if appid is None:
-            return (name,f'{name.title()} not found in start menu.',1,0)
+            return (f'{name.title()} not found in start menu.',1,0)
         
         pid = 0
         if os.path.exists(appid) or "\\" in appid:
             # It's a file path, we can try to get the PID using PassThru
-            command = f'Start-Process "{appid}" -PassThru | Select-Object -ExpandProperty Id'
+            # Escape any single quotes and wrap in single quotes for PowerShell safety
+            safe_appid = appid.replace("'", "''")
+            command = f"Start-Process '{safe_appid}' -PassThru | Select-Object -ExpandProperty Id"
             response, status = self.execute_command(command)
             if status == 0 and response.strip().isdigit():
                 pid = int(response.strip())
         else:
-            # It's an AUMID (Store App)
+            # It's an AUMID (Store App) - validate it only contains expected characters
+            if not appid.replace('\\', '').replace('_', '').replace('.', '').replace('-', '').isalnum():
+                return (f'Invalid app identifier: {appid}', 1, 0)
             command = f'Start-Process "shell:AppsFolder\\{appid}"'
             response, status = self.execute_command(command)
             
         return response, status, pid
     
     def switch_app(self,name:str):
-        apps={app.name:app for app in [self.desktop_state.active_app]+self.desktop_state.apps if app is not None}
-        matched_app:Optional[tuple[str,float]]=process.extractOne(name,list(apps.keys()),score_cutoff=70)
-        if matched_app is None:
+        windows={window.name:window for window in [self.desktop_state.active_window]+self.desktop_state.windows if window is not None}
+        matched_window:Optional[tuple[str,float]]=process.extractOne(name,list(windows.keys()),score_cutoff=70)
+        if matched_window is None:
             return (f'Application {name.title()} not found.',1)
-        app_name,_=matched_app
-        app=apps.get(app_name)
-        target_handle=app.handle
+        window_name,_=matched_window
+        window=windows.get(window_name)
+        target_handle=window.handle
 
         if uia.IsIconic(target_handle):
             uia.ShowWindow(target_handle, win32con.SW_RESTORE)
-            content=f'{app_name.title()} restored from Minimized state.'
+            content=f'{window_name.title()} restored from Minimized state.'
         else:
             self.bring_window_to_top(target_handle)
-            content=f'Switched to {app_name.title()} window.'
+            content=f'Switched to {window_name.title()} window.'
         return content,0
     
     def bring_window_to_top(self, target_handle: int):
@@ -274,6 +298,14 @@ class Desktop:
                 win32gui.ShowWindow(target_handle, win32con.SW_RESTORE)
 
             foreground_handle = win32gui.GetForegroundWindow()
+
+            # Validate both handles before proceeding
+            if not win32gui.IsWindow(foreground_handle):
+                # No valid foreground window, just try to set target as foreground
+                win32gui.SetForegroundWindow(target_handle)
+                win32gui.BringWindowToTop(target_handle)
+                return
+
             foreground_thread, _ = win32process.GetWindowThreadProcessId(foreground_handle)
             target_thread, _ = win32process.GetWindowThreadProcessId(target_handle)
 
@@ -322,7 +354,7 @@ class Desktop:
         x,y=loc
         pg.click(x,y,button=button,clicks=clicks,duration=0.1)
 
-    def type(self,loc:tuple[int,int],text:str,caret_position:Literal['start','end','none']='none',clear:Literal['true','false']='false',press_enter:Literal['true','false']='false'):
+    def type(self,loc:tuple[int,int],text:str,caret_position:Literal['start', 'idle', 'end']='idle',clear:bool|str=False,press_enter:bool|str=False):
         x,y=loc
         pg.leftClick(x,y)
         if caret_position == 'start':
@@ -331,12 +363,16 @@ class Desktop:
             pg.press('end')
         else:
             pass
-        if clear=='true':
+        
+        # Handle both boolean and string 'true'/'false'
+        if clear is True or (isinstance(clear, str) and clear.lower() == 'true'):
             pg.sleep(0.5)
             pg.hotkey('ctrl','a')
             pg.press('backspace')
+            
         pg.typewrite(text,interval=0.02)
-        if press_enter=='true':
+        
+        if press_enter is True or (isinstance(press_enter, str) and press_enter.lower() == 'true'):
             pg.press('enter')
 
     def scroll(self,loc:tuple[int,int]=None,type:Literal['horizontal','vertical']='vertical',direction:Literal['up','down','left','right']='down',wheel_times:int=1)->str|None:
@@ -387,19 +423,20 @@ class Desktop:
         else:
             pg.press(''.join(shortcut))
 
-    def multi_select(self,press_ctrl:Literal['true','false']='false',elements:list[tuple[int,int]|int]=[]):
-        if press_ctrl=='true':
+    def multi_select(self,press_ctrl:bool|str=False,locs:list[tuple[int,int]]=[]):
+        press_ctrl = press_ctrl is True or (isinstance(press_ctrl, str) and press_ctrl.lower() == 'true')
+        if press_ctrl:
             pg.keyDown('ctrl')
-        for element in elements:
-            x,y=element
+        for loc in locs:
+            x,y=loc
             pg.click(x,y,duration=0.2)
             pg.sleep(0.5)
         pg.keyUp('ctrl')
     
-    def multi_edit(self,elements:list[tuple[int,int,str]|tuple[int,str]]):
-        for element in elements:
-            x,y,text=element
-            self.type((x,y),text=text,clear='true')
+    def multi_edit(self,locs:list[tuple[int,int,str]]):
+        for loc in locs:
+            x,y,text=loc
+            self.type((x,y),text=text,clear=True)
     
     def scrape(self,url:str)->str:
         response=requests.get(url,timeout=10)
@@ -407,56 +444,105 @@ class Desktop:
         content=markdownify(html=html)
         return content
     
-    def get_app_from_element(self,element:uia.Control)->App|None:
+    def get_window_from_element(self,element:uia.Control)->Window|None:
         if element is None:
             return None
         top_window=element.GetTopLevelControl()
         if top_window is None:
             return None
         handle=top_window.NativeWindowHandle
-        apps,_=self.get_apps()
-        for app in apps:
-            if app.handle==handle:
-                return app
+        windows,_=self.get_windows()
+        for window in windows:
+            if window.handle==handle:
+                return window
         return None
     
-    def is_app_visible(self,app:uia.Control)->bool:
-        is_minimized=self.get_app_status(app)!=Status.MINIMIZED
-        size=app.BoundingRectangle
+    def is_window_visible(self,window:uia.Control)->bool:
+        is_minimized=self.get_window_status(window)!=Status.MINIMIZED
+        size=window.BoundingRectangle
         area=size.width()*size.height()
-        is_overlay=self.is_overlay_app(app)
+        is_overlay=self.is_overlay_window(window)
         return not is_overlay and is_minimized and area>10
     
-    def is_overlay_app(self,element:uia.Control) -> bool:
+    def is_overlay_window(self,element:uia.Control) -> bool:
         no_children = len(element.GetChildren()) == 0
         is_name = "Overlay" in element.Name.strip()
         return no_children or is_name
 
-    def get_controls_handles(self):
+    def get_controls_handles(self,optimized:bool=False):
         handles = set()
-        root=uia.GetRootControl()
-        children=root.GetChildren()
-        for child in children:
-            handles.add(child.NativeWindowHandle)
+        # For even more faster results (still under development)
+        def callback(hwnd, _):
+            try:
+                # Validate handle before checking properties
+                if win32gui.IsWindow(hwnd) and win32gui.IsWindowVisible(hwnd) and is_window_on_current_desktop(hwnd):
+                    handles.add(hwnd)
+            except Exception:
+                # Skip invalid handles without logging (common during window enumeration)
+                pass
+
+        win32gui.EnumWindows(callback, None)
+
+        if desktop_hwnd:= win32gui.FindWindow('Progman',None):
+            handles.add(desktop_hwnd)
+        if taskbar_hwnd:= win32gui.FindWindow('Shell_TrayWnd',None):
+            handles.add(taskbar_hwnd)
+        if secondary_taskbar_hwnd:= win32gui.FindWindow('Shell_SecondaryTrayWnd',None):
+            handles.add(secondary_taskbar_hwnd)
         return handles
 
-    def get_active_app(self,apps:list[App]|None=None)->App|None:
+    def get_active_window(self,windows:list[Window]|None=None)->Window|None:
         try:
-            if apps is None:
-                apps,_=self.get_apps()
-            handle=uia.GetForegroundWindow()
-            for app in apps:
-                if app.handle!=handle:
+            if windows is None:
+                windows,_=self.get_windows()
+            active_window=self.get_foreground_window()
+            if active_window.ClassName=="Progman":
+                return None
+            active_window_handle=active_window.NativeWindowHandle
+            for window in windows:
+                if window.handle!=active_window_handle:
                     continue
-                return app
+                return window
+            # In case active window is not present in the windows list
+            return Window(**{
+                "name":active_window.Name,
+                "is_browser":self.is_window_browser(active_window),
+                "depth":0,
+                "bounding_box":BoundingBox(
+                    left=active_window.BoundingRectangle.left,
+                    top=active_window.BoundingRectangle.top,
+                    right=active_window.BoundingRectangle.right,
+                    bottom=active_window.BoundingRectangle.bottom,
+                    width=active_window.BoundingRectangle.width(),
+                    height=active_window.BoundingRectangle.height()
+                ),
+                "status":self.get_window_status(active_window),
+                "handle":active_window_handle,
+                "process_id":active_window.ProcessId,
+            })
         except Exception as ex:
-            logger.error(f"Error in get_active_app: {ex}")
+            logger.error(f"Error in get_active_window: {ex}")
         return None
+
+    def get_foreground_window(self)->uia.Control:
+        handle=uia.GetForegroundWindow()
+        active_window=self.get_window_from_element_handle(handle)
+        return active_window
+
+    def get_window_from_element_handle(self, element_handle: int) -> uia.Control:
+        current = uia.ControlFromHandle(element_handle)
+        root_handle = uia.GetRootControl().NativeWindowHandle
         
-    def get_apps(self,controls_handles:set[int]|None=None) -> tuple[list[App],set[int]]:
+        while True:
+            parent = current.GetParentControl()
+            if parent is None or parent.NativeWindowHandle == root_handle:
+                return current
+            current = parent
+        
+    def get_windows(self,controls_handles:set[int]|None=None) -> tuple[list[Window],set[int]]:
         try:
-            apps = []
-            handles = set()
+            windows = []
+            window_handles = set()
             controls_handles=controls_handles or self.get_controls_handles()
             for depth, hwnd in enumerate(controls_handles):
                 try:
@@ -465,7 +551,7 @@ class Desktop:
                     continue
                 
                 # Filter out Overlays (e.g. NVIDIA, Steam)
-                if self.is_overlay_app(child):
+                if self.is_overlay_window(child):
                     continue
 
                 if isinstance(child,(uia.WindowControl,uia.PaneControl)):
@@ -474,15 +560,14 @@ class Desktop:
                         continue
                         
                     if window_pattern.CanMinimize and window_pattern.CanMaximize:
-                        status = self.get_app_status(child)
+                        status = self.get_window_status(child)
                         
                         bounding_rect=child.BoundingRectangle
                         if bounding_rect.isempty() and status!=Status.MINIMIZED:
                             continue
 
-                        apps.append(App(**{
+                        windows.append(Window(**{
                             "name":child.Name,
-                            "runtime_id":tuple(child.GetRuntimeId()),
                             "depth":depth,
                             "status":status,
                             "bounding_box":BoundingBox(
@@ -495,13 +580,13 @@ class Desktop:
                             ),
                             "handle":child.NativeWindowHandle,
                             "process_id":child.ProcessId,
-                            "is_browser":self.is_app_browser(child)
+                            "is_browser":self.is_window_browser(child)
                         }))
-                        handles.add(child.NativeWindowHandle)
+                        window_handles.add(child.NativeWindowHandle)
         except Exception as ex:
-            logger.error(f"Error in get_apps: {ex}")
-            apps = []
-        return apps,handles
+            logger.error(f"Error in get_windows: {ex}")
+            windows = []
+        return windows,window_handles
     
     def get_xpath_from_element(self,element:uia.Control):
         current=element
@@ -556,9 +641,13 @@ class Desktop:
         return "Local Account" if response.strip()=='Local' else "Microsoft Account" if status==0 else "Local Account"
     
     def get_dpi_scaling(self):
-        user32 = ctypes.windll.user32
-        dpi = user32.GetDpiForSystem()
-        return dpi / 96.0
+        try:
+            user32 = ctypes.windll.user32
+            dpi = user32.GetDpiForSystem()
+            return dpi / 96.0 if dpi > 0 else 1.0
+        except Exception:
+            # Fallback to standard DPI if system call fails
+            return 1.0
     
     def get_screen_size(self)->Size:
         width, height = uia.GetVirtualScreenSize()
@@ -568,12 +657,11 @@ class Desktop:
         try:
             return ImageGrab.grab(all_screens=True)
         except Exception as e:
-            logger.warning(f"Failed to capture all screens: {e}. Fallback to primary.")
+            logger.warning(f"Failed to capture virtual screen, using primary screen")
             return pg.screenshot()
 
     def get_annotated_screenshot(self, nodes: list[TreeElementNode]) -> Image.Image:
         screenshot = self.get_screenshot()
-        sleep(0.10)
         # Add padding
         padding = 5
         width = int(screenshot.width + (1.5 * padding))

windows_mcp/desktop/views.py

@@ -6,9 +6,15 @@ from PIL.Image import Image
 from enum import Enum
 
 class Browser(Enum):
-    CHROME='Chrome'
-    EDGE='Edge'
-    FIREFOX='Firefox'
+    CHROME='chrome'
+    EDGE='msedge'
+    FIREFOX='firefox'
+
+    @classmethod
+    def has_process(cls, process_name: str) -> bool:
+        if not hasattr(cls, '_process_names'):
+            cls._process_names = {f'{b.value}.exe' for b in cls}
+        return process_name.lower() in cls._process_names
 
 class Status(Enum):
     MAXIMIZED='Maximized'
@@ -18,9 +24,8 @@ class Status(Enum):
 
 
 @dataclass
-class App:
+class Window:
     name:str
-    runtime_id:tuple[int]
     is_browser:bool
     depth:int
     status:Status
@@ -41,20 +46,32 @@ class Size:
 
 @dataclass
 class DesktopState:
-    apps:list[App]
-    active_app:Optional[App]
+    active_desktop:dict
+    all_desktops:list[dict]
+    active_window:Optional[Window]
+    windows:list[Window]
     screenshot:Optional[Image]=None
     tree_state:Optional[TreeState]=None
 
-    def active_app_to_string(self):
-        if self.active_app is None:
-            return 'No active app found'
+    def active_desktop_to_string(self):
+        desktop_name=self.active_desktop.get('name')
+        headers=["Name"]
+        return tabulate([[desktop_name]], headers=headers, tablefmt="simple")
+
+    def desktops_to_string(self):
+        headers=["Name"]
+        rows=[[desktop.get("name")] for desktop in self.all_desktops]
+        return tabulate(rows, headers=headers, tablefmt="simple")
+
+    def active_window_to_string(self):
+        if not self.active_window:
+            return 'No active window found'
         headers = ["Name", "Depth", "Status", "Width", "Height", "Handle"]
-        return tabulate([self.active_app.to_row()], headers=headers, tablefmt="simple")
+        return tabulate([self.active_window.to_row()], headers=headers, tablefmt="simple")
 
-    def apps_to_string(self):
-        if not self.apps:
-            return 'No apps running in background'
+    def windows_to_string(self):
+        if not self.windows:
+            return 'No windows found'
         headers = ["Name", "Depth", "Status", "Width", "Height", "Handle"]
-        rows = [app.to_row() for app in self.apps]
+        rows = [window.to_row() for window in self.windows]
         return tabulate(rows, headers=headers, tablefmt="simple")