@openhands/extensions 0.4.2 → 0.6.0

package/skills/github-repo-monitor/scripts/main.py

@@ -108,14 +108,47 @@ def fire_callback(
         print(f"Callback error (non-fatal): {exc}")
 
 
-# ── State management ───────────────────────────────────────────────────────────
+# ── State persistence (KV store with local-file fallback) ─────────────────────
 
-def _state_file_path() -> str:
-    """Derive a persistent storage path from WORKSPACE_BASE.
+_KV_TOKEN = os.environ.get("AUTOMATION_KV_TOKEN", "")
+_KV_BASE = os.environ.get("AUTOMATION_API_URL", "").rstrip("/")
+_STATE_KEY = "state"
 
-    WORKSPACE_BASE = {root}/automation-runs/{run_id}
-    State lives two levels up at {root}/automation-state/.
-    """
+
+def _kv_available() -> bool:
+    return bool(_KV_TOKEN and _KV_BASE)
+
+
+def _kv_get(key: str) -> dict | None:
+    req = urllib.request.Request(
+        f"{_KV_BASE}/v1/kv/{key}",
+        headers={"Authorization": f"Bearer {_KV_TOKEN}"},
+    )
+    try:
+        with urllib.request.urlopen(req) as r:
+            return json.loads(r.read())["value"]
+    except urllib.error.HTTPError as exc:
+        if exc.code == 404:
+            return None
+        raise
+
+
+def _kv_set(key: str, value: dict) -> None:
+    req = urllib.request.Request(
+        f"{_KV_BASE}/v1/kv/{key}",
+        data=json.dumps(value).encode(),
+        headers={
+            "Authorization": f"Bearer {_KV_TOKEN}",
+            "Content-Type": "application/json",
+        },
+        method="PUT",
+    )
+    with urllib.request.urlopen(req) as r:
+        r.read()
+
+
+def _state_file_path() -> str:
+    """Derive a persistent storage path from WORKSPACE_BASE (file fallback only)."""
     workspace_base = os.environ.get("WORKSPACE_BASE", "")
     event_payload = json.loads(os.environ.get("AUTOMATION_EVENT_PAYLOAD", "{}"))
     automation_id = event_payload.get("automation_id", "default")
@@ -137,25 +170,42 @@ def _default_since() -> str:
     ).strftime("%Y-%m-%dT%H:%M:%SZ")
 
 
-def load_state(path: str) -> dict:
+def _default_state() -> dict:
+    return {
+        "version": 1,
+        "repo": REPO,
+        "last_poll": _default_since(),
+        "conversations": {},
+        "processed_comment_ids": [],
+    }
+
+
+def load_state() -> dict:
+    if _kv_available():
+        data = _kv_get(_STATE_KEY)
+        if data is not None:
+            print("State loaded from KV store")
+            return data
+        return _default_state()
+    path = _state_file_path()
     if os.path.exists(path):
         try:
             with open(path) as f:
                 return json.load(f)
         except (json.JSONDecodeError, OSError) as exc:
             print(f"Warning: state file {path} unreadable ({exc}); starting fresh")
-    return {
-        "version": 1,
-        "repo": REPO,
-        "last_poll": _default_since(),
-        "conversations": {},       # issue_number (str) → ConversationRecord
-        "processed_comment_ids": [],  # list of int comment IDs already handled
-    }
+    return _default_state()
 
 
-def save_state(path: str, state: dict) -> None:
+def save_state(state: dict) -> None:
+    if _kv_available():
+        _kv_set(_STATE_KEY, state)
+        print("State saved to KV store")
+        return
+    path = _state_file_path()
     with open(path, "w") as f:
         json.dump(state, f, indent=2)
+    print(f"State saved to {path}")
 
 
 # ── GitHub API helpers ─────────────────────────────────────────────────────────
@@ -797,8 +847,7 @@ def _check_conversation_completion(
 
 def main() -> str | None:
     """Run one polling cycle. Returns the last conversation ID created, if any."""
-    state_path = _state_file_path()
-    state = load_state(state_path)
+    state = load_state()
 
     agent_url = os.environ.get("AGENT_SERVER_URL", "").rstrip("/")
     api_key = _get_env_key()
@@ -899,8 +948,7 @@ def main() -> str | None:
     state["processed_comment_ids"] = trimmed
     state["conversations"] = conversations
 
-    save_state(state_path, state)
-    print(f"State saved → {state_path}")
+    save_state(state)
     return last_conversation_id
 
 

package/skills/github-repo-monitor/tests/test_main.py

@@ -29,10 +29,20 @@ def _make_comment(body="hello @openhands", login="octocat", user_type="User"):
 
 # ── State file tests ───────────────────────────────────────────────────────────
 
+def _no_kv(path):
+    """Return a pair of patches: KV disabled, file path pinned to *path*."""
+    return (
+        patch("main._kv_available", return_value=False),
+        patch("main._state_file_path", return_value=path),
+    )
+
+
 class TestLoadState(unittest.TestCase):
 
     def test_missing_file_returns_default(self):
-        state = main.load_state("/nonexistent/path/state.json")
+        with patch("main._kv_available", return_value=False), \
+             patch("main._state_file_path", return_value="/nonexistent/path/state.json"):
+            state = main.load_state()
         self.assertIn("conversations", state)
         self.assertIn("processed_comment_ids", state)
         self.assertEqual(state["version"], 1)
@@ -42,7 +52,9 @@ class TestLoadState(unittest.TestCase):
             json.dump({"version": 1, "custom": "value", "conversations": {}}, f)
             path = f.name
         try:
-            state = main.load_state(path)
+            with patch("main._kv_available", return_value=False), \
+                 patch("main._state_file_path", return_value=path):
+                state = main.load_state()
             self.assertEqual(state["custom"], "value")
         finally:
             os.unlink(path)
@@ -52,7 +64,9 @@ class TestLoadState(unittest.TestCase):
             f.write("{this is not valid json!!!}")
             path = f.name
         try:
-            state = main.load_state(path)
+            with patch("main._kv_available", return_value=False), \
+                 patch("main._state_file_path", return_value=path):
+                state = main.load_state()
             # Should return the default state rather than raising.
             self.assertIn("conversations", state)
             self.assertEqual(state["version"], 1)
@@ -64,7 +78,9 @@ class TestLoadState(unittest.TestCase):
             f.write("")
             path = f.name
         try:
-            state = main.load_state(path)
+            with patch("main._kv_available", return_value=False), \
+                 patch("main._state_file_path", return_value=path):
+                state = main.load_state()
             self.assertIn("conversations", state)
         finally:
             os.unlink(path)
@@ -81,14 +97,55 @@ class TestSaveAndLoadRoundtrip(unittest.TestCase):
         with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
             path = f.name
         try:
-            main.save_state(path, data)
-            loaded = main.load_state(path)
+            with patch("main._kv_available", return_value=False), \
+                 patch("main._state_file_path", return_value=path):
+                main.save_state(data)
+                loaded = main.load_state()
             self.assertEqual(loaded["conversations"]["42"]["conversation_id"], "abc")
             self.assertEqual(loaded["processed_comment_ids"], [1, 2, 3])
         finally:
             os.unlink(path)
 
 
+class TestKVState(unittest.TestCase):
+
+    def test_load_reads_value_from_kv(self):
+        kv_data = {"version": 1, "conversations": {"7": {"conversation_id": "kv-c"}},
+                   "processed_comment_ids": []}
+        with patch("main._kv_available", return_value=True), \
+             patch("main._kv_get", return_value=kv_data) as mock_get:
+            state = main.load_state()
+        self.assertEqual(state["conversations"]["7"]["conversation_id"], "kv-c")
+        mock_get.assert_called_once_with("state")
+
+    def test_load_returns_default_on_kv_miss(self):
+        with patch("main._kv_available", return_value=True), \
+             patch("main._kv_get", return_value=None):
+            state = main.load_state()
+        self.assertIn("conversations", state)
+        self.assertEqual(state["version"], 1)
+
+    def test_save_writes_to_kv(self):
+        data = {"version": 1, "conversations": {}, "processed_comment_ids": []}
+        with patch("main._kv_available", return_value=True), \
+             patch("main._kv_set") as mock_set:
+            main.save_state(data)
+        mock_set.assert_called_once_with("state", data)
+
+    def test_kv_error_propagates_on_load(self):
+        with patch("main._kv_available", return_value=True), \
+             patch("main._kv_get", side_effect=RuntimeError("network error")):
+            with self.assertRaises(RuntimeError):
+                main.load_state()
+
+    def test_kv_error_propagates_on_save(self):
+        data = {"version": 1, "conversations": {}, "processed_comment_ids": []}
+        with patch("main._kv_available", return_value=True), \
+             patch("main._kv_set", side_effect=RuntimeError("network error")):
+            with self.assertRaises(RuntimeError):
+                main.save_state(data)
+
+
 # ── Bot detection tests ────────────────────────────────────────────────────────
 
 class TestIsBotComment(unittest.TestCase):
@@ -214,8 +271,10 @@ class TestProcessedIdDeduplication(unittest.TestCase):
         try:
             state = {"version": 1, "conversations": {},
                      "processed_comment_ids": [101, 202, 303]}
-            main.save_state(path, state)
-            loaded = main.load_state(path)
+            with patch("main._kv_available", return_value=False), \
+                 patch("main._state_file_path", return_value=path):
+                main.save_state(state)
+                loaded = main.load_state()
             self.assertIn(101, loaded["processed_comment_ids"])
             self.assertIn(303, loaded["processed_comment_ids"])
             self.assertNotIn(404, loaded["processed_comment_ids"])