diff --git a/nanobot/heartbeat/service.py b/nanobot/heartbeat/service.py index 3c1a6aa..ab2bfac 100644 --- a/nanobot/heartbeat/service.py +++ b/nanobot/heartbeat/service.py @@ -1,6 +1,7 @@ """Heartbeat service - periodic agent wake-up to check for tasks.""" import asyncio +import re from pathlib import Path from typing import Any, Callable, Coroutine @@ -35,6 +36,15 @@ def _is_heartbeat_empty(content: str | None) -> bool: return True +def _is_heartbeat_ok_response(response: str | None) -> bool: + """Return True only for an exact HEARTBEAT_OK token (with simple markdown wrappers).""" + if not response: + return False + + normalized = re.sub(r"[\s`*_>\-]", "", response).upper() + return normalized == HEARTBEAT_OK_TOKEN.replace("_", "") + + class HeartbeatService: """ Periodic heartbeat service that wakes the agent to check for tasks. @@ -75,6 +85,9 @@ class HeartbeatService: if not self.enabled: logger.info("Heartbeat disabled") return + if self._running: + logger.warning("Heartbeat already running") + return self._running = True self._task = asyncio.create_task(self._run_loop()) @@ -115,7 +128,7 @@ class HeartbeatService: response = await self.on_heartbeat(HEARTBEAT_PROMPT) # Check if agent said "nothing to do" - if HEARTBEAT_OK_TOKEN.replace("_", "") in response.upper().replace("_", ""): + if _is_heartbeat_ok_response(response): logger.info("Heartbeat: OK (no action needed)") else: logger.info("Heartbeat: completed task") diff --git a/tests/test_heartbeat_service.py b/tests/test_heartbeat_service.py new file mode 100644 index 0000000..52d1b96 --- /dev/null +++ b/tests/test_heartbeat_service.py @@ -0,0 +1,40 @@ +import asyncio + +import pytest + +from nanobot.heartbeat.service import ( + HeartbeatService, + _is_heartbeat_ok_response, +) + + +def test_heartbeat_ok_response_requires_exact_token() -> None: + assert _is_heartbeat_ok_response("HEARTBEAT_OK") + assert _is_heartbeat_ok_response("`HEARTBEAT_OK`") + assert _is_heartbeat_ok_response("**HEARTBEAT_OK**") + + assert not _is_heartbeat_ok_response("HEARTBEAT_OK, done") + assert not _is_heartbeat_ok_response("done HEARTBEAT_OK") + assert not _is_heartbeat_ok_response("HEARTBEAT_NOT_OK") + + +@pytest.mark.asyncio +async def test_start_is_idempotent(tmp_path) -> None: + async def _on_heartbeat(_: str) -> str: + return "HEARTBEAT_OK" + + service = HeartbeatService( + workspace=tmp_path, + on_heartbeat=_on_heartbeat, + interval_s=9999, + enabled=True, + ) + + await service.start() + first_task = service._task + await service.start() + + assert service._task is first_task + + service.stop() + await asyncio.sleep(0)