diff --git a/aw_qt/manager.py b/aw_qt/manager.py index 2aa05a6..ad6d40a 100644 --- a/aw_qt/manager.py +++ b/aw_qt/manager.py @@ -200,9 +200,12 @@ def stop(self) -> None: self.started = False def toggle(self, testing: bool) -> None: - if self.started: + if self.is_alive(): self.stop() else: + if self.started: + # Process died unexpectedly, clean up state + self.stop() self.start(testing) def is_alive(self) -> bool: diff --git a/aw_qt/trayicon.py b/aw_qt/trayicon.py index e3d5dfd..da42281 100644 --- a/aw_qt/trayicon.py +++ b/aw_qt/trayicon.py @@ -6,7 +6,7 @@ import time import webbrowser from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional import aw_core from PyQt6 import QtCore @@ -73,6 +73,9 @@ def open_dir(d: str) -> None: class TrayIcon(QSystemTrayIcon): + MAX_AUTO_RESTARTS = 3 + RESTART_WINDOW_SECONDS = 600 # 10 minutes + def __init__( self, manager: Manager, @@ -86,12 +89,30 @@ def __init__( self.manager = manager self.testing = testing + self._restart_timestamps: Dict[str, List[float]] = {} self.root_url = f"http://localhost:{5666 if self.testing else 5600}" self.activated.connect(self.on_activated) self._build_rootmenu() + def _recent_restart_count(self, module_name: str) -> int: + """Count restarts within the sliding time window.""" + now = time.monotonic() + timestamps = self._restart_timestamps.get(module_name, []) + cutoff = now - self.RESTART_WINDOW_SECONDS + return sum(1 for t in timestamps if t > cutoff) + + def _record_restart(self, module_name: str) -> None: + """Record a restart timestamp and prune old entries.""" + now = time.monotonic() + cutoff = now - self.RESTART_WINDOW_SECONDS + timestamps = self._restart_timestamps.get(module_name, []) + # Prune old timestamps and add current + self._restart_timestamps[module_name] = [ + t for t in timestamps if t > cutoff + ] + [now] + def on_activated(self, reason: QSystemTrayIcon.ActivationReason) -> None: if reason == QSystemTrayIcon.ActivationReason.DoubleClick: open_webui(self.root_url) @@ -137,11 +158,25 @@ def _build_rootmenu(self) -> None: def show_module_failed_dialog(module: Module) -> None: box = QMessageBox(self._parent) box.setIcon(QMessageBox.Icon.Warning) - box.setText(f"Module {module.name} quit unexpectedly") + recent = self._recent_restart_count(module.name) + box.setText( + f"Module {module.name} quit unexpectedly" + + ( + f" after {recent} auto-restart attempts" + f" in {self.RESTART_WINDOW_SECONDS // 60} minutes" + if recent >= self.MAX_AUTO_RESTARTS + else "" + ) + ) box.setDetailedText(module.read_log(self.testing)) restart_button = QPushButton("Restart", box) - restart_button.clicked.connect(module.start) + + def on_manual_restart() -> None: + self._restart_timestamps.pop(module.name, None) + module.start(self.testing) + + restart_button.clicked.connect(on_manual_restart) box.addButton(restart_button, QMessageBox.ButtonRole.AcceptRole) box.setStandardButtons(QMessageBox.StandardButton.Cancel) @@ -153,31 +188,54 @@ def rebuild_modules_menu() -> None: module: Module = action.data() alive = module.is_alive() action.setChecked(alive) - # print(module.text(), alive) - # TODO: Do it in a better way, singleShot isn't pretty... QtCore.QTimer.singleShot(2000, rebuild_modules_menu) QtCore.QTimer.singleShot(2000, rebuild_modules_menu) def check_module_status() -> None: unexpected_exits = self.manager.get_unexpected_stops() - if unexpected_exits: - for module in unexpected_exits: + for module in unexpected_exits: + recent = self._recent_restart_count(module.name) + if recent < self.MAX_AUTO_RESTARTS: + logger.info( + f"Auto-restarting crashed module {module.name} " + f"(attempt {recent + 1}/{self.MAX_AUTO_RESTARTS}" + f" in {self.RESTART_WINDOW_SECONDS // 60}min window)" + ) + module.stop() # Clean up state + module.start(self.testing) + self._record_restart(module.name) + self.showMessage( + "ActivityWatch", + f"Module {module.name} crashed and was auto-restarted", + QSystemTrayIcon.MessageIcon.Warning, + 5000, + ) + else: + logger.warning( + f"Module {module.name} exceeded max auto-restarts " + f"({self.MAX_AUTO_RESTARTS})" + ) show_module_failed_dialog(module) module.stop() - # TODO: Do it in a better way, singleShot isn't pretty... - QtCore.QTimer.singleShot(2000, rebuild_modules_menu) + QtCore.QTimer.singleShot(5000, check_module_status) - QtCore.QTimer.singleShot(2000, check_module_status) + QtCore.QTimer.singleShot(5000, check_module_status) def _build_modulemenu(self, moduleMenu: QMenu) -> None: moduleMenu.clear() def add_module_menuitem(module: Module) -> None: title = module.name - ac = moduleMenu.addAction(title, lambda: module.toggle(self.testing)) + + def on_toggle(m: Module = module) -> None: + m.toggle(self.testing) + # Reset auto-restart timestamps on manual toggle + self._restart_timestamps.pop(m.name, None) + + ac = moduleMenu.addAction(title, on_toggle) ac.setData(module) ac.setCheckable(True) diff --git a/tests/test_manager.py b/tests/test_manager.py new file mode 100644 index 0000000..0eb6379 --- /dev/null +++ b/tests/test_manager.py @@ -0,0 +1,121 @@ +"""Unit tests for the Module manager.""" + +import subprocess +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from aw_qt.manager import Module + + +@pytest.fixture +def module(): + """Create a test module with a mock path.""" + return Module("aw-test-module", Path("/usr/bin/true"), "system") + + +class TestModuleToggle: + """Tests for Module.toggle() behavior, especially with crashed processes.""" + + def test_toggle_starts_stopped_module(self, module): + """Toggle should start a module that was never started.""" + assert not module.started + assert not module.is_alive() + + with patch.object(module, "start") as mock_start: + module.toggle(testing=True) + mock_start.assert_called_once_with(True) + + def test_toggle_stops_running_module(self, module): + """Toggle should stop a module that is running.""" + # Simulate a running process + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.returncode = None # None means still running + + # After terminate+wait, process should report as dead + def fake_wait(): + mock_proc.returncode = -15 # SIGTERM + + mock_proc.wait.side_effect = fake_wait + module._process = mock_proc + module.started = True + + module.toggle(testing=True) + assert not module.started + assert module._process is None + + def test_toggle_restarts_crashed_module(self, module): + """Toggle should restart a module that crashed (started=True, process dead). + + This is the key fix: previously, toggle checked `self.started` instead of + `self.is_alive()`, so the first click would only call stop() (cleaning up + state) without starting, requiring a second click to actually start. + """ + # Simulate a crashed process: started=True but process exited + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.returncode = 1 # Non-zero = exited + module._process = mock_proc + module.started = True + + assert not module.is_alive() # Process is dead + assert module.started # But flag says started + + with patch.object(module, "start") as mock_start: + module.toggle(testing=True) + # Should have cleaned up and started in one toggle + mock_start.assert_called_once_with(True) + assert not module.started # stop() was called to clean up + + +class TestModuleIsAlive: + """Tests for Module.is_alive() behavior.""" + + def test_is_alive_no_process(self, module): + assert not module.is_alive() + + def test_is_alive_running(self, module): + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.returncode = None + module._process = mock_proc + assert module.is_alive() + + def test_is_alive_exited(self, module): + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.returncode = 0 + module._process = mock_proc + assert not module.is_alive() + + +class TestGetUnexpectedStops: + """Tests for Manager.get_unexpected_stops().""" + + def test_detects_crashed_module(self): + from aw_qt.manager import Manager + + with patch.object(Manager, "discover_modules"): + mgr = Manager(testing=True) + + mod = Module("aw-test", Path("/usr/bin/true"), "system") + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.returncode = 1 + mod._process = mock_proc + mod.started = True + mgr.modules = [mod] + + unexpected = mgr.get_unexpected_stops() + assert len(unexpected) == 1 + assert unexpected[0].name == "aw-test" + + def test_ignores_intentionally_stopped(self): + from aw_qt.manager import Manager + + with patch.object(Manager, "discover_modules"): + mgr = Manager(testing=True) + + mod = Module("aw-test", Path("/usr/bin/true"), "system") + mod.started = False # Intentionally stopped + mgr.modules = [mod] + + unexpected = mgr.get_unexpected_stops() + assert len(unexpected) == 0