diff --git a/libs/computer/computer/computer.py b/libs/computer/computer/computer.py index 191c611d..4249f3f4 100644 --- a/libs/computer/computer/computer.py +++ b/libs/computer/computer/computer.py @@ -181,24 +181,25 @@ class Computer: self.logger.debug("Telemetry disabled - skipping initialization tracking") async def __aenter__(self): - """Enter async context manager.""" + """Start the computer.""" + await self.run() return self async def __aexit__(self, exc_type, exc_val, exc_tb): - """Exit async context manager.""" - pass + """Stop the computer.""" + await self.disconnect() def __enter__(self): - """Enter synchronous context manager.""" - # Run the event loop to call the async run method + """Start the computer.""" + # Run the event loop to call the async enter method loop = asyncio.get_event_loop() - loop.run_until_complete(self.run()) + loop.run_until_complete(self.__aenter__()) return self def __exit__(self, exc_type, exc_val, exc_tb): - """Exit synchronous context manager.""" - # We could add cleanup here if needed in the future - pass + """Stop the computer.""" + loop = asyncio.get_event_loop() + loop.run_until_complete(self.__aexit__(exc_type, exc_val, exc_tb)) async def run(self) -> Optional[str]: """Initialize the VM and computer interface.""" @@ -473,9 +474,14 @@ class Computer: duration_ms = (time.time() - start_time) * 1000 self.logger.debug(f"Computer initialization took {duration_ms:.2f}ms") return + + async def disconnect(self) -> None: + """Disconnect from the computer's WebSocket interface.""" + if self._interface: + self._interface.close() async def stop(self) -> None: - """Stop computer control.""" + """Disconnect from the computer's WebSocket interface and stop the computer.""" start_time = time.time() try: @@ -496,6 +502,7 @@ class Computer: await self.config.vm_provider.__aexit__(None, None, None) self._provider_context = None + await self.disconnect() self.logger.info("Computer stopped") except Exception as e: self.logger.debug(f"Error during cleanup: {e}") # Log as debug since this might be expected diff --git a/libs/computer/computer/providers/cloud/provider.py b/libs/computer/computer/providers/cloud/provider.py index bc2226a8..1cfba161 100644 --- a/libs/computer/computer/providers/cloud/provider.py +++ b/libs/computer/computer/providers/cloud/provider.py @@ -52,11 +52,11 @@ class CloudProvider(BaseVMProvider): return [] async def run_vm(self, image: str, name: str, run_opts: Dict[str, Any], storage: Optional[str] = None) -> Dict[str, Any]: - logger.warning("CloudProvider.run_vm is not implemented") - return {"name": name, "status": "unavailable", "message": "CloudProvider is not implemented"} + # logger.warning("CloudProvider.run_vm is not implemented") + return {"name": name, "status": "unavailable", "message": "CloudProvider.run_vm is not implemented"} async def stop_vm(self, name: str, storage: Optional[str] = None) -> Dict[str, Any]: - logger.warning("CloudProvider.stop_vm is not implemented") + logger.warning("CloudProvider.stop_vm is not implemented. To clean up resources, please use Computer.disconnect()") return {"name": name, "status": "stopped", "message": "CloudProvider is not implemented"} async def update_vm(self, name: str, update_opts: Dict[str, Any], storage: Optional[str] = None) -> Dict[str, Any]: