implement review remarks

Author: mlikasam-askui

src/askui/tools/askui/askui_controller.py

@@ -225,7 +225,11 @@ def connect(self) -> None:
         self._start_execution()
         self.set_display(self._display)
         if self._settings.clean_virtual_displays:
+            logger.info(
+                "clean_virtual_displays is enabled. Removing all virtual displays ... "
+            )
             self.remove_virtual_displays()
+            logger.info("Virtual displays removed.")
 
     def _get_stub(self) -> controller_v1.ControllerAPIStub:
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
@@ -1395,14 +1399,6 @@ def remove_virtual_displays(self) -> None:
         self._send_command(command)
         self._reporter.add_message("AgentOS", "remove_virtual_displays() -> done")
 
-        logger.debug(
-            (
-                "Setting display to 1 to ensure that the "
-                "controller is using the real display"
-            )
-        )
-        self.set_display(1)
-
     @staticmethod
     def _decode_file_payload(base64_data: str) -> Image.Image | str:
         try:

src/askui/tools/askui/askui_ui_controller_grpc/generated/AgentOS_Send_Request_2501.py

@@ -406,7 +406,7 @@ class GetFileCommand(BaseModel):
 
 class RemoveVirtualDisplaysCommand(BaseModel):
     name: Literal['RemoveVirtualDisplays'] = 'RemoveVirtualDisplays'
-    parameters: List[Any] = []
+    parameters: List[str] = Field(default=[], max_length=0)
 
 Command =Union[
         GetSystemInfoCommand,

src/askui/tools/store/computer/experimental/get_file.py

@@ -49,6 +49,7 @@ def __init__(self, agent_os: AgentOs | None = None) -> None:
             agent_os=agent_os,
             required_tags=[ToolTags.SCALED_AGENT_OS.value],
         )
+        self.is_cacheable = True
 
     def __call__(self, absolute_file_path: str) -> Image.Image | str:
         return self.agent_os.get_file(absolute_file_path)

src/askui/tools/store/computer/experimental/remove_virtual_displays.py

@@ -34,14 +34,11 @@ def __init__(self, agent_os: AgentOs | None = None) -> None:
                 "add_window_as_virtual_display_tool) to restore a normal setup before "
                 "continuing automation."
             ),
-            input_schema={
-                "type": "object",
-                "properties": {},
-                "required": [],
-            },
             agent_os=agent_os,
         )
 
+        self.is_cacheable = True
+
     def __call__(self) -> str:
         self.agent_os.remove_virtual_displays()
         return "Removed virtual displays; only physical displays remain active."