1  """ByteRover memory plugin — MemoryProvider interface.
  2  
  3  Persistent memory via the ByteRover CLI (``brv``). Organizes knowledge into
  4  a hierarchical context tree with tiered retrieval (fuzzy text → LLM-driven
  5  search). Local-first with optional cloud sync.
  6  
  7  Original PR #3499 by hieuntg81, adapted to MemoryProvider ABC.
  8  
  9  Requires: ``brv`` CLI installed (npm install -g byterover-cli or
 10  curl -fsSL https://byterover.dev/install.sh | sh).
 11  
 12  Config via environment variables (profile-scoped via each profile's .env):
 13    BRV_API_KEY   — ByteRover API key (for cloud features, optional for local)
 14  
 15  Working directory: $HERMES_HOME/byterover/ (profile-scoped context tree)
 16  """
 17  
 18  from __future__ import annotations
 19  
 20  import json
 21  import logging
 22  import os
 23  import shutil
 24  import subprocess
 25  import threading
 26  from pathlib import Path
 27  from typing import Any, Dict, List, Optional
 28  
 29  from agent.memory_provider import MemoryProvider
 30  from tools.registry import tool_error
 31  
 32  logger = logging.getLogger(__name__)
 33  
 34  # Timeouts
 35  _QUERY_TIMEOUT = 10   # brv query — should be fast
 36  _CURATE_TIMEOUT = 120  # brv curate — may involve LLM processing
 37  
 38  # Minimum lengths to filter noise
 39  _MIN_QUERY_LEN = 10
 40  _MIN_OUTPUT_LEN = 20
 41  
 42  
 43  # ---------------------------------------------------------------------------
 44  # brv binary resolution (cached, thread-safe)
 45  # ---------------------------------------------------------------------------
 46  
 47  _brv_path_lock = threading.Lock()
 48  _cached_brv_path: Optional[str] = None
 49  
 50  
 51  def _resolve_brv_path() -> Optional[str]:
 52      """Find the brv binary on PATH or well-known install locations."""
 53      global _cached_brv_path
 54      with _brv_path_lock:
 55          if _cached_brv_path is not None:
 56              return _cached_brv_path if _cached_brv_path != "" else None
 57  
 58      found = shutil.which("brv")
 59      if not found:
 60          home = Path.home()
 61          candidates = [
 62              home / ".brv-cli" / "bin" / "brv",
 63              Path("/usr/local/bin/brv"),
 64              home / ".npm-global" / "bin" / "brv",
 65          ]
 66          for c in candidates:
 67              if c.exists():
 68                  found = str(c)
 69                  break
 70  
 71      with _brv_path_lock:
 72          if _cached_brv_path is not None:
 73              return _cached_brv_path if _cached_brv_path != "" else None
 74          _cached_brv_path = found or ""
 75      return found
 76  
 77  
 78  def _run_brv(args: List[str], timeout: int = _QUERY_TIMEOUT,
 79               cwd: str = None) -> dict:
 80      """Run a brv CLI command. Returns {success, output, error}."""
 81      brv_path = _resolve_brv_path()
 82      if not brv_path:
 83          return {"success": False, "error": "brv CLI not found. Install: npm install -g byterover-cli"}
 84  
 85      cmd = [brv_path] + args
 86      effective_cwd = cwd or str(_get_brv_cwd())
 87      Path(effective_cwd).mkdir(parents=True, exist_ok=True)
 88  
 89      env = os.environ.copy()
 90      brv_bin_dir = str(Path(brv_path).parent)
 91      env["PATH"] = brv_bin_dir + os.pathsep + env.get("PATH", "")
 92  
 93      try:
 94          result = subprocess.run(
 95              cmd, capture_output=True, text=True,
 96              timeout=timeout, cwd=effective_cwd, env=env,
 97          )
 98          stdout = result.stdout.strip()
 99          stderr = result.stderr.strip()
100  
101          if result.returncode == 0:
102              return {"success": True, "output": stdout}
103          return {"success": False, "error": stderr or stdout or f"brv exited {result.returncode}"}
104  
105      except subprocess.TimeoutExpired:
106          return {"success": False, "error": f"brv timed out after {timeout}s"}
107      except FileNotFoundError:
108          global _cached_brv_path
109          with _brv_path_lock:
110              _cached_brv_path = None
111          return {"success": False, "error": "brv CLI not found"}
112      except Exception as e:
113          return {"success": False, "error": str(e)}
114  
115  
116  def _get_brv_cwd() -> Path:
117      """Profile-scoped working directory for the brv context tree."""
118      from hermes_constants import get_hermes_home
119      return get_hermes_home() / "byterover"
120  
121  
122  # ---------------------------------------------------------------------------
123  # Tool schemas
124  # ---------------------------------------------------------------------------
125  
126  QUERY_SCHEMA = {
127      "name": "brv_query",
128      "description": (
129          "Search ByteRover's persistent knowledge tree for relevant context. "
130          "Returns memories, project knowledge, architectural decisions, and "
131          "patterns from previous sessions. Use for any question where past "
132          "context would help."
133      ),
134      "parameters": {
135          "type": "object",
136          "properties": {
137              "query": {"type": "string", "description": "What to search for."},
138          },
139          "required": ["query"],
140      },
141  }
142  
143  CURATE_SCHEMA = {
144      "name": "brv_curate",
145      "description": (
146          "Store important information in ByteRover's persistent knowledge tree. "
147          "Use for architectural decisions, bug fixes, user preferences, project "
148          "patterns — anything worth remembering across sessions. ByteRover's LLM "
149          "automatically categorizes and organizes the memory."
150      ),
151      "parameters": {
152          "type": "object",
153          "properties": {
154              "content": {"type": "string", "description": "The information to remember."},
155          },
156          "required": ["content"],
157      },
158  }
159  
160  STATUS_SCHEMA = {
161      "name": "brv_status",
162      "description": "Check ByteRover status — CLI version, context tree stats, cloud sync state.",
163      "parameters": {"type": "object", "properties": {}, "required": []},
164  }
165  
166  
167  # ---------------------------------------------------------------------------
168  # MemoryProvider implementation
169  # ---------------------------------------------------------------------------
170  
171  class ByteRoverMemoryProvider(MemoryProvider):
172      """ByteRover persistent memory via the brv CLI."""
173  
174      def __init__(self):
175          self._cwd = ""
176          self._session_id = ""
177          self._turn_count = 0
178          self._sync_thread: Optional[threading.Thread] = None
179  
180      @property
181      def name(self) -> str:
182          return "byterover"
183  
184      def is_available(self) -> bool:
185          """Check if brv CLI is installed. No network calls."""
186          return _resolve_brv_path() is not None
187  
188      def get_config_schema(self):
189          return [
190              {
191                  "key": "api_key",
192                  "description": "ByteRover API key (optional, for cloud sync)",
193                  "secret": True,
194                  "env_var": "BRV_API_KEY",
195                  "url": "https://app.byterover.dev",
196              },
197          ]
198  
199      def initialize(self, session_id: str, **kwargs) -> None:
200          self._cwd = str(_get_brv_cwd())
201          self._session_id = session_id
202          self._turn_count = 0
203          Path(self._cwd).mkdir(parents=True, exist_ok=True)
204  
205      def system_prompt_block(self) -> str:
206          if not _resolve_brv_path():
207              return ""
208          return (
209              "# ByteRover Memory\n"
210              "Active. Persistent knowledge tree with hierarchical context.\n"
211              "Use brv_query to search past knowledge, brv_curate to store "
212              "important facts, brv_status to check state."
213          )
214  
215      def prefetch(self, query: str, *, session_id: str = "") -> str:
216          """Run brv query synchronously before the agent's first LLM call.
217  
218          Blocks until the query completes (up to _QUERY_TIMEOUT seconds), ensuring
219          the result is available as context before the model is called.
220          """
221          if not query or len(query.strip()) < _MIN_QUERY_LEN:
222              return ""
223          result = _run_brv(
224              ["query", "--", query.strip()[:5000]],
225              timeout=_QUERY_TIMEOUT, cwd=self._cwd,
226          )
227          if result["success"] and result.get("output"):
228              output = result["output"].strip()
229              if len(output) > _MIN_OUTPUT_LEN:
230                  return f"## ByteRover Context\n{output}"
231          return ""
232  
233      def queue_prefetch(self, query: str, *, session_id: str = "") -> None:
234          """No-op: prefetch() now runs synchronously at turn start."""
235          pass
236  
237      def sync_turn(self, user_content: str, assistant_content: str, *, session_id: str = "") -> None:
238          """Curate the conversation turn in background (non-blocking)."""
239          self._turn_count += 1
240  
241          # Only curate substantive turns
242          if len(user_content.strip()) < _MIN_QUERY_LEN:
243              return
244  
245          def _sync():
246              try:
247                  combined = f"User: {user_content[:2000]}\nAssistant: {assistant_content[:2000]}"
248                  _run_brv(
249                      ["curate", "--", combined],
250                      timeout=_CURATE_TIMEOUT, cwd=self._cwd,
251                  )
252              except Exception as e:
253                  logger.debug("ByteRover sync failed: %s", e)
254  
255          # Wait for previous sync
256          if self._sync_thread and self._sync_thread.is_alive():
257              self._sync_thread.join(timeout=5.0)
258  
259          self._sync_thread = threading.Thread(
260              target=_sync, daemon=True, name="brv-sync"
261          )
262          self._sync_thread.start()
263  
264      def on_memory_write(self, action: str, target: str, content: str) -> None:
265          """Mirror built-in memory writes to ByteRover."""
266          if action not in ("add", "replace") or not content:
267              return
268  
269          def _write():
270              try:
271                  label = "User profile" if target == "user" else "Agent memory"
272                  _run_brv(
273                      ["curate", "--", f"[{label}] {content}"],
274                      timeout=_CURATE_TIMEOUT, cwd=self._cwd,
275                  )
276              except Exception as e:
277                  logger.debug("ByteRover memory mirror failed: %s", e)
278  
279          t = threading.Thread(target=_write, daemon=True, name="brv-memwrite")
280          t.start()
281  
282      def on_pre_compress(self, messages: List[Dict[str, Any]]) -> str:
283          """Extract insights before context compression discards turns."""
284          if not messages:
285              return ""
286  
287          # Build a summary of messages about to be compressed
288          parts = []
289          for msg in messages[-10:]:  # last 10 messages
290              role = msg.get("role", "")
291              content = msg.get("content", "")
292              if isinstance(content, str) and content.strip() and role in ("user", "assistant"):
293                  parts.append(f"{role}: {content[:500]}")
294  
295          if not parts:
296              return ""
297  
298          combined = "\n".join(parts)
299  
300          def _flush():
301              try:
302                  _run_brv(
303                      ["curate", "--", f"[Pre-compression context]\n{combined}"],
304                      timeout=_CURATE_TIMEOUT, cwd=self._cwd,
305                  )
306                  logger.info("ByteRover pre-compression flush: %d messages", len(parts))
307              except Exception as e:
308                  logger.debug("ByteRover pre-compression flush failed: %s", e)
309  
310          t = threading.Thread(target=_flush, daemon=True, name="brv-flush")
311          t.start()
312          return ""
313  
314      def get_tool_schemas(self) -> List[Dict[str, Any]]:
315          return [QUERY_SCHEMA, CURATE_SCHEMA, STATUS_SCHEMA]
316  
317      def handle_tool_call(self, tool_name: str, args: dict, **kwargs) -> str:
318          if tool_name == "brv_query":
319              return self._tool_query(args)
320          elif tool_name == "brv_curate":
321              return self._tool_curate(args)
322          elif tool_name == "brv_status":
323              return self._tool_status()
324          return tool_error(f"Unknown tool: {tool_name}")
325  
326      def shutdown(self) -> None:
327          if self._sync_thread and self._sync_thread.is_alive():
328              self._sync_thread.join(timeout=10.0)
329  
330      # -- Tool implementations ------------------------------------------------
331  
332      def _tool_query(self, args: dict) -> str:
333          query = args.get("query", "")
334          if not query:
335              return tool_error("query is required")
336  
337          result = _run_brv(
338              ["query", "--", query.strip()[:5000]],
339              timeout=_QUERY_TIMEOUT, cwd=self._cwd,
340          )
341  
342          if not result["success"]:
343              return tool_error(result.get("error", "Query failed"))
344  
345          output = result.get("output", "").strip()
346          if not output or len(output) < _MIN_OUTPUT_LEN:
347              return json.dumps({"result": "No relevant memories found."})
348  
349          # Truncate very long results
350          if len(output) > 8000:
351              output = output[:8000] + "\n\n[... truncated]"
352  
353          return json.dumps({"result": output})
354  
355      def _tool_curate(self, args: dict) -> str:
356          content = args.get("content", "")
357          if not content:
358              return tool_error("content is required")
359  
360          result = _run_brv(
361              ["curate", "--", content],
362              timeout=_CURATE_TIMEOUT, cwd=self._cwd,
363          )
364  
365          if not result["success"]:
366              return tool_error(result.get("error", "Curate failed"))
367  
368          return json.dumps({"result": "Memory curated successfully."})
369  
370      def _tool_status(self) -> str:
371          result = _run_brv(["status"], timeout=15, cwd=self._cwd)
372          if not result["success"]:
373              return tool_error(result.get("error", "Status check failed"))
374          return json.dumps({"status": result.get("output", "")})
375  
376  
377  # ---------------------------------------------------------------------------
378  # Plugin entry point
379  # ---------------------------------------------------------------------------
380  
381  def register(ctx) -> None:
382      """Register ByteRover as a memory provider plugin."""
383      ctx.register_memory_provider(ByteRoverMemoryProvider())