#!/usr/bin/env python3 from __future__ import annotations import json import os import sys import urllib.error import urllib.parse import urllib.request from pathlib import Path from typing import Any PLUGIN_ROOT = Path(__file__).resolve().parents[1] if str(PLUGIN_ROOT) not in sys.path: sys.path.insert(0, str(PLUGIN_ROOT)) from memory_gateway_plugin.output import dumps_safe, short_id, summarize_data USER_ID = "test_user_memory_gateway_plugin" AGENT_ID = "test_hermes_memory_gateway_plugin" WORKSPACE_ID = "test_workspace_memory_gateway_plugin" SESSION_ID = "test_session_memory_gateway_plugin_001" def _assert_test_user(user_id: str) -> None: if not user_id.startswith("test_user_"): raise ValueError("cleanup_refuses_non_test_user") def _gateway_url() -> str: return os.environ.get("MEMORY_GATEWAY_URL", "http://127.0.0.1:1934").rstrip("/") def _api_key() -> str: return os.environ.get("MEMORY_GATEWAY_API_KEY", "") def _request(method: str, endpoint: str, payload: dict[str, Any] | None = None) -> dict[str, Any]: headers = {"Content-Type": "application/json"} if _api_key(): headers["X-API-Key"] = _api_key() body = None if payload is None else json.dumps(payload, ensure_ascii=False).encode("utf-8") req = urllib.request.Request(_gateway_url() + endpoint, data=body, headers=headers, method=method) try: with urllib.request.urlopen(req, timeout=15) as response: raw = response.read().decode("utf-8") return {"ok": True, "status_code": getattr(response, "status", 200), "data": json.loads(raw) if raw else {}} except urllib.error.HTTPError as exc: return {"ok": False, "status_code": exc.code, "error": str(exc.reason)[:300]} except Exception as exc: return {"ok": False, "status_code": None, "error": str(exc)[:300]} def _search_candidates() -> dict[str, Any]: return _request( "POST", "/v1/memory/search", { "query": "integration_test plugin safe_to_delete Memory Gateway plugin", "user_id": USER_ID, "agent_id": AGENT_ID, "workspace_id": WORKSPACE_ID, "session_id": SESSION_ID, "tags": ["integration_test"], "limit": 100, }, ) def _audit_candidate_ids() -> list[str]: result = _request("GET", "/v1/audit?limit=1000") if not result.get("ok"): return [] ids: list[str] = [] for row in result.get("data") or []: if row.get("actor_user_id") != USER_ID: continue if row.get("actor_agent_id") not in {AGENT_ID, None, ""}: continue if row.get("target_type") == "memory" and row.get("action") in {"upsert_memory", "feedback:incorrect", "feedback:duplicate", "feedback:outdated"}: target_id = row.get("target_id") if target_id and target_id not in ids: ids.append(target_id) return ids def _memory_from_result(item: dict[str, Any]) -> dict[str, Any] | None: memory = item.get("memory") if isinstance(memory, dict) and memory.get("id"): return memory return None def _is_cleanup_candidate(memory: dict[str, Any]) -> bool: if memory.get("user_id") != USER_ID: return False tags = set(memory.get("tags") or []) return bool(tags.intersection({"integration_test", "safe_to_delete", "plugin"})) def _delete_memory(memory_id: str) -> dict[str, Any]: query = urllib.parse.urlencode({"user_id": USER_ID, "agent_id": AGENT_ID, "workspace_id": WORKSPACE_ID, "session_id": SESSION_ID}) return _request("DELETE", f"/v1/memory/{urllib.parse.quote(memory_id)}?{query}") def _feedback_memory(memory_id: str) -> dict[str, Any]: return _request( "POST", f"/v1/memory/{urllib.parse.quote(memory_id)}/feedback", { "user_id": USER_ID, "agent_id": AGENT_ID, "workspace_id": WORKSPACE_ID, "session_id": SESSION_ID, "feedback": "incorrect", "comment": "cleanup marker for integration test memory", }, ) def run(user_id: str = USER_ID) -> dict[str, Any]: _assert_test_user(user_id) if user_id != USER_ID: return {"ok": False, "error": "script_is_scoped_to_fixed_test_user", "user_id": user_id} search = _search_candidates() if not search.get("ok"): return {"ok": False, "search": {"ok": False, "status_code": search.get("status_code"), "error": search.get("error")}, "deleted": 0, "feedback_marked": 0, "skipped": 0} rows = (search.get("data") or {}).get("results") or [] memory_ids = _audit_candidate_ids() deleted = 0 feedback_marked = 0 skipped = 0 unable: list[dict[str, Any]] = [] touched: list[str] = [] for item in rows: memory = _memory_from_result(item) if not memory or not _is_cleanup_candidate(memory): skipped += 1 continue memory_id = memory["id"] if memory_id not in memory_ids: memory_ids.append(memory_id) for memory_id in memory_ids: deletion = _delete_memory(memory_id) if deletion.get("ok"): deleted += 1 touched.append(short_id(memory_id)) continue if deletion.get("status_code") == 404: skipped += 1 continue feedback = _feedback_memory(memory_id) if feedback.get("ok"): feedback_marked += 1 touched.append(short_id(memory_id)) else: unable.append({"memory_id": short_id(memory_id), "delete_status": deletion.get("status_code"), "feedback_status": feedback.get("status_code"), "reason": feedback.get("error") or deletion.get("error")}) return { "ok": not unable, "search": {"ok": True, "status_code": search.get("status_code"), "data": summarize_data(search.get("data"))}, "deleted": deleted, "feedback_marked": feedback_marked, "skipped": skipped, "unable_count": len(unable), "unable": unable, "touched_memory_ids": touched, "limitation": "search API returns local MemoryRecord rows plus OpenViking context; cleanup only deletes local MemoryRecord rows for the fixed test user.", } def main() -> int: try: result = run(os.environ.get("MEMORY_GATEWAY_CLEANUP_USER_ID", USER_ID)) except ValueError as exc: result = {"ok": False, "error": str(exc), "deleted": 0, "feedback_marked": 0, "skipped": 0} print(dumps_safe(result)) return 0 if result.get("ok") else 1 if __name__ == "__main__": sys.exit(main())