From 47f5ec484724dd8cab0ab99627d72997197e8ce7 Mon Sep 17 00:00:00 2001 From: Nyx Date: Fri, 8 May 2026 10:52:25 +0530 Subject: [PATCH] fix: bound memory capture and recall output --- agent/memory_manager.py | 11 ++- hermes_cli/config.py | 4 + plugins/memory/supermemory/__init__.py | 45 ++++++++--- .../memory/test_supermemory_provider.py | 75 ++++++++++++++++++- tests/tools/test_session_search.py | 33 ++++++++ tools/memory_search_caps.py | 68 +++++++++++++++++ tools/session_search_tool.py | 22 ++++-- 7 files changed, 238 insertions(+), 20 deletions(-) create mode 100644 tools/memory_search_caps.py diff --git a/agent/memory_manager.py b/agent/memory_manager.py index 1319681d3b196..e1ce70ad5011d 100644 --- a/agent/memory_manager.py +++ b/agent/memory_manager.py @@ -34,6 +34,7 @@ from tools.registry import tool_error logger = logging.getLogger(__name__) +_MEMORY_SEARCH_TOOL_RE = re.compile(r"(search|query|recall|profile)$") # --------------------------------------------------------------------------- @@ -365,7 +366,15 @@ def handle_tool_call( if provider is None: return tool_error(f"No memory provider handles tool '{tool_name}'") try: - return provider.handle_tool_call(tool_name, args, **kwargs) + result = provider.handle_tool_call(tool_name, args, **kwargs) + if _MEMORY_SEARCH_TOOL_RE.search(tool_name): + try: + from tools.memory_search_caps import cap_memory_search_result + + return cap_memory_search_result(result) + except Exception: + logger.debug("Memory search result cap failed for %s", tool_name, exc_info=True) + return result except Exception as e: logger.error( "Memory provider '%s' handle_tool_call(%s) failed: %s", diff --git a/hermes_cli/config.py b/hermes_cli/config.py index cb6753864f157..65f7dd810595d 100644 --- a/hermes_cli/config.py +++ b/hermes_cli/config.py @@ -982,6 +982,10 @@ def _ensure_hermes_home_managed(home: Path): "user_profile_enabled": True, "memory_char_limit": 2200, # ~800 tokens at 2.75 chars/token "user_char_limit": 1375, # ~500 tokens at 2.75 chars/token + # Hard cap for memory/recall tool outputs (session_search, memory-provider + # search/profile/recall tools). Tool boundary enforcement appends a + # visible truncation notice instead of silently flooding context. + "search_result_char_limit": 10000, # External memory provider plugin (empty = built-in only). # Set to a provider name to activate: "openviking", "mem0", # "hindsight", "holographic", "retaindb", "byterover". diff --git a/plugins/memory/supermemory/__init__.py b/plugins/memory/supermemory/__init__.py index f0cbfd60276da..322e69094085a 100644 --- a/plugins/memory/supermemory/__init__.py +++ b/plugins/memory/supermemory/__init__.py @@ -26,6 +26,7 @@ _DEFAULT_MAX_RECALL_RESULTS = 10 _DEFAULT_PROFILE_FREQUENCY = 50 _DEFAULT_CAPTURE_MODE = "all" +_VALID_CAPTURE_MODES = {"all", "everything", "explicit"} _DEFAULT_SEARCH_MODE = "hybrid" _VALID_SEARCH_MODES = ("hybrid", "memories", "documents") _DEFAULT_API_TIMEOUT = 5.0 @@ -120,7 +121,8 @@ def _load_supermemory_config(hermes_home: str) -> dict: config["profile_frequency"] = max(1, min(500, int(config.get("profile_frequency", _DEFAULT_PROFILE_FREQUENCY)))) except Exception: config["profile_frequency"] = _DEFAULT_PROFILE_FREQUENCY - config["capture_mode"] = "everything" if config.get("capture_mode") == "everything" else "all" + raw_capture_mode = str(config.get("capture_mode", _DEFAULT_CAPTURE_MODE)).strip().lower() + config["capture_mode"] = raw_capture_mode if raw_capture_mode in _VALID_CAPTURE_MODES else _DEFAULT_CAPTURE_MODE raw_search_mode = str(config.get("search_mode", _DEFAULT_SEARCH_MODE)).strip().lower() config["search_mode"] = raw_search_mode if raw_search_mode in _VALID_SEARCH_MODES else _DEFAULT_SEARCH_MODE config["entity_context"] = _clamp_entity_context(str(config.get("entity_context", _DEFAULT_ENTITY_CONTEXT))) @@ -439,6 +441,8 @@ def __init__(self): self._entity_context = _DEFAULT_ENTITY_CONTEXT self._api_timeout = _DEFAULT_API_TIMEOUT self._hermes_home = "" + self._platform = "" + self._agent_context = "" self._write_enabled = True self._active = False # Multi-container support @@ -508,6 +512,8 @@ def initialize(self, session_id: str, **kwargs) -> None: self._allowed_containers = [self._container_tag] + list(self._custom_containers) agent_context = kwargs.get("agent_context", "") + self._platform = str(kwargs.get("platform") or "").strip() + self._agent_context = str(agent_context or "").strip() self._write_enabled = agent_context not in ("cron", "flush", "subagent") self._active = bool(self._api_key) self._client = None @@ -563,6 +569,8 @@ def prefetch(self, query: str, *, session_id: str = "") -> str: def sync_turn(self, user_content: str, assistant_content: str, *, session_id: str = "") -> None: if not self._active or not self._auto_capture or not self._write_enabled or not self._client: return + if self._capture_mode == "explicit": + return clean_user = _clean_text_for_capture(user_content) clean_assistant = _clean_text_for_capture(assistant_content) @@ -593,7 +601,7 @@ def _run(): self._sync_thread.start() def on_session_end(self, messages: List[Dict[str, Any]]) -> None: - if not self._active or not self._write_enabled or not self._client or not self._session_id: + if not self._active or not self._auto_capture or not self._write_enabled or not self._client or not self._session_id: return cleaned = [] for message in messages or []: @@ -624,7 +632,14 @@ def _run(): try: self._client.add_memory( content.strip(), - metadata={"source": "hermes_memory", "target": target, "type": "explicit_memory"}, + metadata={ + "source": "hermes_memory", + "target": target, + "type": "explicit_memory", + "saved_at": datetime.now(timezone.utc).isoformat(), + **({"platform": self._platform} if self._platform else {}), + **({"origin": self._agent_context} if self._agent_context else {}), + }, entity_context=self._entity_context, ) except Exception: @@ -692,6 +707,11 @@ def _tool_store(self, args: dict) -> str: metadata = {} metadata.setdefault("type", _detect_category(content)) metadata["source"] = "hermes_tool" + metadata.setdefault("saved_at", datetime.now(timezone.utc).isoformat()) + if self._platform: + metadata.setdefault("platform", self._platform) + if self._agent_context: + metadata.setdefault("origin", self._agent_context) try: result = self._client.add_memory(content, metadata=metadata, entity_context=self._entity_context, container_tag=tag) preview = content[:80] + ("..." if len(content) > 80 else "") @@ -757,15 +777,22 @@ def _tool_profile(self, args: dict) -> str: return tool_error(str(exc)) try: profile = self._client.get_profile(query=query, container_tag=tag) + static_facts = profile["static"] or [] + dynamic_facts = profile["dynamic"] or [] + static_visible = static_facts[:self._max_recall_results] + dynamic_visible = dynamic_facts[:self._max_recall_results] sections = [] - if profile["static"]: - sections.append("## User Profile (Persistent)\n" + "\n".join(f"- {item}" for item in profile["static"])) - if profile["dynamic"]: - sections.append("## Recent Context\n" + "\n".join(f"- {item}" for item in profile["dynamic"])) + if static_visible: + sections.append("## User Profile (Persistent)\n" + "\n".join(f"- {item}" for item in static_visible)) + if dynamic_visible: + sections.append("## Recent Context\n" + "\n".join(f"- {item}" for item in dynamic_visible)) resp: dict[str, Any] = { "profile": "\n\n".join(sections), - "static_count": len(profile["static"]), - "dynamic_count": len(profile["dynamic"]), + "static_count": len(static_facts), + "dynamic_count": len(dynamic_facts), + "visible_static_count": len(static_visible), + "visible_dynamic_count": len(dynamic_visible), + "truncated": len(static_facts) > len(static_visible) or len(dynamic_facts) > len(dynamic_visible), } if tag: resp["container_tag"] = tag diff --git a/tests/plugins/memory/test_supermemory_provider.py b/tests/plugins/memory/test_supermemory_provider.py index 0aee459757f4a..934d6f5816870 100644 --- a/tests/plugins/memory/test_supermemory_provider.py +++ b/tests/plugins/memory/test_supermemory_provider.py @@ -84,12 +84,13 @@ def fake_import(name, *args, **kwargs): def test_load_and_save_config_round_trip(tmp_path): - _save_supermemory_config({"container_tag": "demo-tag", "auto_capture": False}, str(tmp_path)) + _save_supermemory_config({"container_tag": "demo-tag", "auto_capture": False, "capture_mode": "explicit"}, str(tmp_path)) cfg = _load_supermemory_config(str(tmp_path)) # container_tag is kept raw — sanitization happens in initialize() after template resolution assert cfg["container_tag"] == "demo-tag" assert cfg["auto_capture"] is False assert cfg["auto_recall"] is True + assert cfg["capture_mode"] == "explicit" def test_clean_text_for_capture_strips_injected_context(): @@ -153,6 +154,16 @@ def test_sync_turn_persists_cleaned_exchange(provider): assert "[role: assistant]" in content +def test_sync_turn_skips_when_capture_mode_explicit(provider): + provider._capture_mode = "explicit" + provider.sync_turn( + "Please remember this request in long-term memory", + "Absolutely, I will keep that in long-term memory.", + session_id="session-1", + ) + assert provider._client.add_calls == [] + + def test_on_session_end_ingests_clean_messages(provider): messages = [ {"role": "system", "content": "skip"}, @@ -174,7 +185,12 @@ def test_on_memory_write_tracks_thread(provider): assert provider._write_thread is not None provider._write_thread.join(timeout=1) assert len(provider._client.add_calls) == 1 - assert provider._client.add_calls[0]["metadata"]["type"] == "explicit_memory" + metadata = provider._client.add_calls[0]["metadata"] + assert metadata["type"] == "explicit_memory" + assert metadata["source"] == "hermes_memory" + assert metadata["target"] == "memory" + assert metadata["platform"] == "cli" + assert "saved_at" in metadata def test_shutdown_joins_and_clears_threads(provider, monkeypatch): @@ -220,6 +236,11 @@ def test_store_tool_returns_saved_payload(provider): result = json.loads(provider.handle_tool_call("supermemory_store", {"content": "Jordan likes concise docs"})) assert result["saved"] is True assert result["id"] == "mem_123" + metadata = provider._client.add_calls[0]["metadata"] + assert metadata["source"] == "hermes_tool" + assert metadata["type"] == "preference" + assert metadata["platform"] == "cli" + assert "saved_at" in metadata def test_search_tool_formats_results(provider): @@ -253,9 +274,31 @@ def test_profile_tool_formats_sections(provider): result = json.loads(provider.handle_tool_call("supermemory_profile", {})) assert result["static_count"] == 1 assert result["dynamic_count"] == 1 + assert result["visible_static_count"] == 1 + assert result["visible_dynamic_count"] == 1 + assert result["truncated"] is False assert "User Profile (Persistent)" in result["profile"] +def test_profile_tool_limits_visible_profile_context(provider): + provider._max_recall_results = 2 + provider._client.profile_response = { + "static": ["static 1", "static 2", "static 3"], + "dynamic": ["dynamic 1", "dynamic 2", "dynamic 3"], + "search_results": [], + } + result = json.loads(provider.handle_tool_call("supermemory_profile", {})) + assert result["static_count"] == 3 + assert result["dynamic_count"] == 3 + assert result["visible_static_count"] == 2 + assert result["visible_dynamic_count"] == 2 + assert result["truncated"] is True + assert "static 1" in result["profile"] + assert "static 3" not in result["profile"] + assert "dynamic 1" in result["profile"] + assert "dynamic 3" not in result["profile"] + + def test_handle_tool_call_returns_error_when_unconfigured(monkeypatch): monkeypatch.delenv("SUPERMEMORY_API_KEY", raising=False) p = SupermemoryMemoryProvider() @@ -409,3 +452,31 @@ def test_get_config_schema_minimal(): assert len(schema) == 1 assert schema[0]["key"] == "api_key" assert schema[0]["secret"] is True + + +def test_memory_manager_caps_memory_search_tool_results(monkeypatch): + from agent.memory_manager import MemoryManager + + class Provider: + name = "dummy" + + def is_available(self): + return True + + def initialize(self, session_id, **kwargs): + pass + + def get_tool_schemas(self): + return [{"name": "dummy_search", "description": "", "parameters": {}}] + + def handle_tool_call(self, tool_name, args, **kwargs): + return "x" * 12000 + + monkeypatch.setattr("tools.memory_search_caps.get_memory_search_result_char_limit", lambda: 10000) + manager = MemoryManager() + manager.add_provider(Provider()) + + result = manager.handle_tool_call("dummy_search", {}) + + assert len(result) <= 10000 + assert "[Result truncated at 10K chars." in result diff --git a/tests/tools/test_session_search.py b/tests/tools/test_session_search.py index 468a492ad8e00..e8df9691f0ff0 100644 --- a/tests/tools/test_session_search.py +++ b/tests/tools/test_session_search.py @@ -5,6 +5,7 @@ import time import pytest +from tools.memory_search_caps import cap_memory_search_result, format_memory_search_truncation_notice from tools.session_search_tool import ( _format_timestamp, _format_conversation, @@ -14,6 +15,7 @@ _HIDDEN_SESSION_SOURCES, MAX_SESSION_CHARS, SESSION_SEARCH_SCHEMA, + session_search, ) @@ -184,6 +186,37 @@ def test_multiword_window_maximises_coverage(self): assert result.lower().count("alpha beta") == 2 +class TestMemorySearchResultCap: + def test_cap_appends_visible_notice_and_respects_limit(self): + result = cap_memory_search_result("x" * 12000, limit=10000) + assert len(result) <= 10000 + assert result.endswith(format_memory_search_truncation_notice(10000)) + assert "Result truncated at 10K chars" in result + + def test_session_search_caps_large_result_at_tool_boundary(self, monkeypatch): + from unittest.mock import MagicMock + + async def fake_summarize(*_args, **_kwargs): + return "s" * 12000 + + monkeypatch.setattr("tools.session_search_tool._summarize_session", fake_summarize) + monkeypatch.setattr("model_tools._run_async", lambda coro: asyncio.run(coro)) + monkeypatch.setattr("tools.memory_search_caps.get_memory_search_result_char_limit", lambda: 10000) + + mock_db = MagicMock() + mock_db.search_messages.return_value = [ + {"session_id": "s1", "source": "cli", "session_started": 1709500000, "model": "test"}, + ] + mock_db.get_session.return_value = {"id": "s1", "parent_session_id": None, "source": "cli", "started_at": 1709500000} + mock_db.get_messages_as_conversation.return_value = [{"role": "user", "content": "message"}] + + result = session_search(query="message", db=mock_db, limit=1) + + assert len(result) <= 10000 + assert result.endswith(format_memory_search_truncation_notice(10000)) + assert "Result truncated at 10K chars" in result + + class TestSessionSearchConcurrency: def test_defaults_to_three(self): assert _get_session_search_max_concurrency() == 3 diff --git a/tools/memory_search_caps.py b/tools/memory_search_caps.py new file mode 100644 index 0000000000000..d20e52c19c6aa --- /dev/null +++ b/tools/memory_search_caps.py @@ -0,0 +1,68 @@ +"""Shared cap for raw memory-search style tool outputs.""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + +DEFAULT_MEMORY_SEARCH_RESULT_CHAR_LIMIT = 10_000 +MIN_MEMORY_SEARCH_RESULT_CHAR_LIMIT = 1_000 +MAX_MEMORY_SEARCH_RESULT_CHAR_LIMIT = 1_000_000 + + +def get_memory_search_result_char_limit(default: int = DEFAULT_MEMORY_SEARCH_RESULT_CHAR_LIMIT) -> int: + """Read memory.search_result_char_limit with sane bounds.""" + try: + from hermes_cli.config import load_config + + config = load_config() + except Exception: + return default + + memory_config = config.get("memory", {}) if isinstance(config, dict) else {} + raw = memory_config.get("search_result_char_limit", default) if isinstance(memory_config, dict) else default + try: + value = int(raw) + except (TypeError, ValueError): + return default + return max(MIN_MEMORY_SEARCH_RESULT_CHAR_LIMIT, min(value, MAX_MEMORY_SEARCH_RESULT_CHAR_LIMIT)) + + +def format_memory_search_truncation_notice(limit: int) -> str: + """Return the visible truncation notice appended to capped outputs.""" + if limit % 1000 == 0: + label = f"{limit // 1000}K" + else: + label = f"{limit}" + return ( + f"[Result truncated at {label} chars. Use a more specific query, narrower date range, " + "or recall: filter to get focused results.]" + ) + + +def cap_memory_search_result(raw: str, *, limit: int | None = None) -> str: + """Cap a raw memory/search tool result and append a visible truncation notice. + + The final returned string stays at or below the configured limit, so callers + cannot accidentally inject oversized recall payloads into the agent context. + """ + if raw is None: + raw = "" + if not isinstance(raw, str): + raw = str(raw) + + effective_limit = get_memory_search_result_char_limit() if limit is None else int(limit) + effective_limit = max(MIN_MEMORY_SEARCH_RESULT_CHAR_LIMIT, min(effective_limit, MAX_MEMORY_SEARCH_RESULT_CHAR_LIMIT)) + if len(raw) <= effective_limit: + return raw + + notice = format_memory_search_truncation_notice(effective_limit) + suffix = "\n" + notice + keep = max(0, effective_limit - len(suffix)) + logger.info( + "Truncated memory/search tool result from %d to %d chars", + len(raw), + effective_limit, + ) + return raw[:keep].rstrip() + suffix diff --git a/tools/session_search_tool.py b/tools/session_search_tool.py index efc450b322ea0..18c4efcdd4050 100644 --- a/tools/session_search_tool.py +++ b/tools/session_search_tool.py @@ -25,6 +25,7 @@ from typing import Dict, Any, List, Optional, Union from agent.auxiliary_client import async_call_llm, extract_content_or_reasoning +from tools.memory_search_caps import cap_memory_search_result, get_memory_search_result_char_limit MAX_SESSION_CHARS = 100_000 MAX_SUMMARY_TOKENS = 10000 @@ -265,6 +266,10 @@ async def _summarize_session( _HIDDEN_SESSION_SOURCES = ("tool",) +def _cap_result(payload: str) -> str: + return cap_memory_search_result(payload) + + def _list_recent_sessions(db, limit: int, current_session_id: str = None) -> str: """Return metadata for the most recent sessions (no LLM calls).""" try: @@ -310,13 +315,13 @@ def _list_recent_sessions(db, limit: int, current_session_id: str = None) -> str if len(results) >= limit: break - return json.dumps({ + return _cap_result(json.dumps({ "success": True, "mode": "recent", "results": results, "count": len(results), "message": f"Showing {len(results)} most recent sessions. Use a keyword query to search specific topics.", - }, ensure_ascii=False) + }, ensure_ascii=False)) except Exception as e: logging.error("Error listing recent sessions: %s", e, exc_info=True) return tool_error(f"Failed to list recent sessions: {e}", success=False) @@ -372,13 +377,13 @@ def session_search( ) if not raw_results: - return json.dumps({ + return _cap_result(json.dumps({ "success": True, "query": query, "results": [], "count": 0, "message": "No matching sessions found.", - }, ensure_ascii=False) + }, ensure_ascii=False)) # Resolve child sessions to their parent — delegation stores detailed # content in child sessions, but the user's conversation is the parent. @@ -480,10 +485,10 @@ async def _bounded_summary(text: str, meta: Dict[str, Any]) -> Optional[str]: "Session summarization timed out after 60 seconds", exc_info=True, ) - return json.dumps({ + return _cap_result(json.dumps({ "success": False, "error": "Session summarization timed out. Try a more specific query or reduce the limit.", - }, ensure_ascii=False) + }, ensure_ascii=False)) summaries = [] for (session_id, match_info, conversation_text, session_meta), result in zip(tasks, results): @@ -518,13 +523,13 @@ async def _bounded_summary(text: str, meta: Dict[str, Any]) -> Optional[str]: summaries.append(entry) - return json.dumps({ + return _cap_result(json.dumps({ "success": True, "query": query, "results": summaries, "count": len(summaries), "sessions_searched": len(seen_sessions), - }, ensure_ascii=False) + }, ensure_ascii=False)) except Exception as e: logging.error("Session search failed: %s", e, exc_info=True) @@ -602,4 +607,5 @@ def check_session_search_requirements() -> bool: current_session_id=kw.get("current_session_id")), check_fn=check_session_search_requirements, emoji="🔍", + max_result_size_chars=get_memory_search_result_char_limit(), )