Gateway error handling

pyproject.toml

@@ -24,7 +24,7 @@ prompt-toolkit = "^3.0.43"
 rich = "^13.7.1"
 platformdirs = "^4.2.0"
 pyyaml = "^6.0.2"
-mcp = "1.8.0"
+mcp = "1.11.0"
 
 ## POS dependencies
 urllib3 = ">= 1.25.3"

src/pieces/_vendor/pieces_os_client/wrapper/websockets/base_websocket.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional,TYPE_CHECKING, List
+from typing import Callable, Optional,TYPE_CHECKING, List, Self
 import websocket
 import threading
 from abc import ABC, abstractmethod
@@ -149,7 +149,7 @@ def is_running(cls) -> bool:
 		return False
 
 	@classmethod
-	def get_instance(cls) -> Optional[type]:
+	def get_instance(cls) -> Optional[Self]:
 		"""
 		Get the singleton instance of the class.
 

src/pieces/command_interface/open_command.py

@@ -54,6 +54,12 @@ def add_arguments(self, parser: argparse.ArgumentParser):
             action="store_true",
             help="Opens Pieces Settings",
         )
+        parser.add_argument(
+            "--ltm",
+            dest="ltm",
+            action="store_true",
+            help="Opens Pieces LTM (Long-Term Memory)",
+        )
 
     def execute(self, **kwargs) -> int:
         """Execute the open command."""

src/pieces/command_interface/simple_commands.py

@@ -1,4 +1,5 @@
 import argparse
+from urllib3.exceptions import MaxRetryError
 from pieces.base_command import BaseCommand
 from pieces.headless.models.base import CommandResult
 from pieces.headless.models.version import create_version_success
@@ -205,3 +206,39 @@ def execute(self, **kwargs) -> CommandResult:
         return CommandResult(
             0, create_version_success(__version__, Settings.pieces_os_version)
         )
+
+
+class RestartPiecesOSCommand(BaseCommand):
+    """Command to restart PiecesOS."""
+
+    def get_name(self) -> str:
+        return "restart"
+
+    def get_help(self) -> str:
+        return "Restart PiecesOS"
+
+    def get_description(self) -> str:
+        return "Restart the PiecesOS"
+
+    def get_examples(self) -> list[str]:
+        return ["pieces restart"]
+
+    def get_docs(self) -> str:
+        return URLs.CLI_RESTART_DOCS.value
+
+    def add_arguments(self, parser: argparse.ArgumentParser):
+        pass
+
+    def execute(self, **kwargs) -> int:
+        try:
+            Settings.pieces_client.os_api.os_restart()
+        except MaxRetryError:
+            pass
+        if Settings.pieces_client.is_pieces_running(15):
+            Settings.logger.print("[green]PiecesOS restarted successfully.")
+            return 0
+        else:
+            Settings.logger.print(
+                "[red]Failed to restart PiecesOS. Please run `pieces open`."
+            )
+        return 1

src/pieces/copilot/ltm.py

@@ -115,7 +115,7 @@ def capture(application):
     application.exit(result=s if s.dimensions else None)
 
 
-def check_ltm(docs=None) -> bool:
+def check_ltm(docs=None, auto_enable=False) -> bool:
     # Update the local cache
     Settings.pieces_client.copilot.context.ltm.ltm_status = Settings.pieces_client.work_stream_pattern_engine_api.workstream_pattern_engine_processors_vision_status()
     if Settings.pieces_client.copilot.context.ltm.is_enabled:
@@ -124,7 +124,7 @@ def check_ltm(docs=None) -> bool:
     if Settings.headless_mode:
         raise HeadlessLTMNotEnabledError()
 
-    if not Settings.logger.confirm(
+    if not auto_enable and not Settings.logger.confirm(
         "Pieces LTM must be running, do you want to enable it?",
     ):
         return False
@@ -277,8 +277,8 @@ def _open_ltm():
         Settings.show_error(f"Error in enabling the LTM: {e}")
 
 
-def enable_ltm():
-    if check_ltm():
+def enable_ltm(auto_enable: bool = False) -> bool:
+    if check_ltm(None, auto_enable):
         # window = add_qrcodes()  # TODO: Clean at exist
         # if not window:
         #     Settings.show_error(

src/pieces/core/open_command.py

@@ -1,20 +1,26 @@
 from pieces.urls import URLs
 from pieces.settings import Settings
+from pieces.copilot.ltm import enable_ltm
 
 def open_command(**kwargs):
     from pieces._vendor.pieces_os_client.models.inactive_os_server_applet import InactiveOSServerApplet
     from pieces._vendor.pieces_os_client.models.os_applet_enum import OSAppletEnum
     copilot = kwargs.get("copilot", False)
     drive = kwargs.get("drive", False,)
     settings = kwargs.get("settings", False)
+    ltm = kwargs.get("ltm", False)
 
     # Let's try to Open POS
     health = Settings.pieces_client.open_pieces_os()
 
-    if (drive or copilot or settings) and not health:
+    if (drive or copilot or settings or ltm) and not health:
         Settings.logger.print("PiecesOS is not running")
         return
 
+    if ltm and enable_ltm(auto_enable=True):
+        Settings.logger.print("[green]LTM is enabled and running[/green]")
+
+
     if copilot:
         URLs.open_website(
             "localhost:"