From 40d338e9667c2d097b4bc8fe5dd0df17b5d7e377 Mon Sep 17 00:00:00 2001 From: helix4u <4317663+helix4u@users.noreply.github.com> Date: Fri, 24 Apr 2026 13:34:27 -0600 Subject: [PATCH] fix(memory): add write origin metadata --- agent/memory_manager.py | 45 ++++++++++++++++- agent/memory_provider.py | 15 ++++-- plugins/memory/byterover/__init__.py | 8 ++- plugins/memory/holographic/__init__.py | 10 +++- plugins/memory/honcho/__init__.py | 8 ++- plugins/memory/openviking/__init__.py | 8 ++- plugins/memory/retaindb/__init__.py | 8 ++- plugins/memory/supermemory/__init__.py | 17 ++++++- run_agent.py | 52 ++++++++++++++++++++ tests/agent/test_memory_provider.py | 52 ++++++++++++++++++++ tests/run_agent/test_flush_memories_codex.py | 25 ++++++++++ 11 files changed, 235 insertions(+), 13 deletions(-) diff --git a/agent/memory_manager.py b/agent/memory_manager.py index 2435c3f248399..62cbd6ae1ad5e 100644 --- a/agent/memory_manager.py +++ b/agent/memory_manager.py @@ -31,6 +31,7 @@ import json import logging import re +import inspect from typing import Any, Dict, List, Optional from agent.memory_provider import MemoryProvider @@ -312,7 +313,39 @@ def on_pre_compress(self, messages: List[Dict[str, Any]]) -> str: ) return "\n\n".join(parts) - def on_memory_write(self, action: str, target: str, content: str) -> None: + @staticmethod + def _provider_memory_write_metadata_mode(provider: MemoryProvider) -> str: + """Return how to pass metadata to a provider's memory-write hook.""" + try: + signature = inspect.signature(provider.on_memory_write) + except (TypeError, ValueError): + return "keyword" + + params = list(signature.parameters.values()) + if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params): + return "keyword" + if "metadata" in signature.parameters: + return "keyword" + + accepted = [ + p for p in params + if p.kind in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + ] + if len(accepted) >= 4: + return "positional" + return "legacy" + + def on_memory_write( + self, + action: str, + target: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: """Notify external providers when the built-in memory tool writes. Skips the builtin provider itself (it's the source of the write). @@ -321,7 +354,15 @@ def on_memory_write(self, action: str, target: str, content: str) -> None: if provider.name == "builtin": continue try: - provider.on_memory_write(action, target, content) + metadata_mode = self._provider_memory_write_metadata_mode(provider) + if metadata_mode == "keyword": + provider.on_memory_write( + action, target, content, metadata=dict(metadata or {}) + ) + elif metadata_mode == "positional": + provider.on_memory_write(action, target, content, dict(metadata or {})) + else: + provider.on_memory_write(action, target, content) except Exception as e: logger.debug( "Memory provider '%s' on_memory_write failed: %s", diff --git a/agent/memory_provider.py b/agent/memory_provider.py index 24593e33457c0..535338f4ee2ec 100644 --- a/agent/memory_provider.py +++ b/agent/memory_provider.py @@ -26,7 +26,7 @@ on_turn_start(turn, message, **kwargs) — per-turn tick with runtime context on_session_end(messages) — end-of-session extraction on_pre_compress(messages) -> str — extract before context compression - on_memory_write(action, target, content) — mirror built-in memory writes + on_memory_write(action, target, content, metadata=None) — mirror built-in memory writes on_delegation(task, result, **kwargs) — parent-side observation of subagent work """ @@ -34,7 +34,7 @@ import logging from abc import ABC, abstractmethod -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional logger = logging.getLogger(__name__) @@ -220,12 +220,21 @@ def save_config(self, values: Dict[str, Any], hermes_home: str) -> None: should all have ``env_var`` set and this method stays no-op). """ - def on_memory_write(self, action: str, target: str, content: str) -> None: + def on_memory_write( + self, + action: str, + target: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: """Called when the built-in memory tool writes an entry. action: 'add', 'replace', or 'remove' target: 'memory' or 'user' content: the entry content + metadata: structured provenance for the write, when available. Common + keys include ``write_origin``, ``execution_context``, ``session_id``, + ``parent_session_id``, ``platform``, and ``tool_name``. Use to mirror built-in memory writes to your backend. """ diff --git a/plugins/memory/byterover/__init__.py b/plugins/memory/byterover/__init__.py index 1870e9ab865e2..cb1fd785fbd2b 100644 --- a/plugins/memory/byterover/__init__.py +++ b/plugins/memory/byterover/__init__.py @@ -261,7 +261,13 @@ def _sync(): ) self._sync_thread.start() - def on_memory_write(self, action: str, target: str, content: str) -> None: + def on_memory_write( + self, + action: str, + target: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: """Mirror built-in memory writes to ByteRover.""" if action not in ("add", "replace") or not content: return diff --git a/plugins/memory/holographic/__init__.py b/plugins/memory/holographic/__init__.py index cd4ef07b44c70..e1bcec694ca56 100644 --- a/plugins/memory/holographic/__init__.py +++ b/plugins/memory/holographic/__init__.py @@ -20,7 +20,7 @@ import json import logging import re -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from agent.memory_provider import MemoryProvider from tools.registry import tool_error @@ -240,7 +240,13 @@ def on_session_end(self, messages: List[Dict[str, Any]]) -> None: return self._auto_extract_facts(messages) - def on_memory_write(self, action: str, target: str, content: str) -> None: + def on_memory_write( + self, + action: str, + target: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: """Mirror built-in memory writes as facts.""" if action == "add" and self._store and content: try: diff --git a/plugins/memory/honcho/__init__.py b/plugins/memory/honcho/__init__.py index 6ca32c1dcbb5c..57c0893b70e15 100644 --- a/plugins/memory/honcho/__init__.py +++ b/plugins/memory/honcho/__init__.py @@ -1087,7 +1087,13 @@ def _sync(): ) self._sync_thread.start() - def on_memory_write(self, action: str, target: str, content: str) -> None: + def on_memory_write( + self, + action: str, + target: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: """Mirror built-in user profile writes as Honcho conclusions.""" if action != "add" or target != "user" or not content: return diff --git a/plugins/memory/openviking/__init__.py b/plugins/memory/openviking/__init__.py index 86d7ad5efb1f4..b381e94a76f46 100644 --- a/plugins/memory/openviking/__init__.py +++ b/plugins/memory/openviking/__init__.py @@ -469,7 +469,13 @@ def on_session_end(self, messages: List[Dict[str, Any]]) -> None: except Exception as e: logger.warning("OpenViking session commit failed: %s", e) - def on_memory_write(self, action: str, target: str, content: str) -> None: + def on_memory_write( + self, + action: str, + target: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: """Mirror built-in memory writes to OpenViking as explicit memories.""" if not self._client or action != "add" or not content: return diff --git a/plugins/memory/retaindb/__init__.py b/plugins/memory/retaindb/__init__.py index 62121410d41c7..07b4e622f563c 100644 --- a/plugins/memory/retaindb/__init__.py +++ b/plugins/memory/retaindb/__init__.py @@ -744,7 +744,13 @@ def _dispatch(self, tool_name: str, args: dict) -> Any: # ── Optional hooks ───────────────────────────────────────────────────── - def on_memory_write(self, action: str, target: str, content: str) -> None: + def on_memory_write( + self, + action: str, + target: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: """Mirror built-in memory writes to RetainDB.""" if action != "add" or not content or not self._client: return diff --git a/plugins/memory/supermemory/__init__.py b/plugins/memory/supermemory/__init__.py index f0cbfd60276da..2e2735e1315ad 100644 --- a/plugins/memory/supermemory/__init__.py +++ b/plugins/memory/supermemory/__init__.py @@ -614,17 +614,30 @@ def on_session_end(self, messages: List[Dict[str, Any]]) -> None: except Exception: logger.warning("Supermemory session ingest failed", exc_info=True) - def on_memory_write(self, action: str, target: str, content: str) -> None: + def on_memory_write( + self, + action: str, + target: str, + content: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: if not self._active or not self._write_enabled or not self._client: return if action != "add" or not (content or "").strip(): return + write_metadata = { + "source": "hermes_memory", + "target": target, + "type": "explicit_memory", + } + if metadata: + write_metadata.update(metadata) def _run(): try: self._client.add_memory( content.strip(), - metadata={"source": "hermes_memory", "target": target, "type": "explicit_memory"}, + metadata=write_metadata, entity_context=self._entity_context, ) except Exception: diff --git a/run_agent.py b/run_agent.py index 4911e4899e385..18956bf566f9d 100644 --- a/run_agent.py +++ b/run_agent.py @@ -1437,6 +1437,8 @@ def __init__( # Track conversation messages for session logging self._session_messages: List[Dict[str, Any]] = [] + self._memory_write_origin = "assistant_tool" + self._memory_write_context = "foreground" # Cached system prompt -- built once per session, only rebuilt on compression self._cached_system_prompt: Optional[str] = None @@ -3047,7 +3049,10 @@ def _run_review(): quiet_mode=True, platform=self.platform, provider=self.provider, + parent_session_id=self.session_id, ) + review_agent._memory_write_origin = "background_review" + review_agent._memory_write_context = "background_review" review_agent._memory_store = self._memory_store review_agent._memory_enabled = self._memory_enabled review_agent._user_profile_enabled = self._user_profile_enabled @@ -3095,6 +3100,32 @@ def _run_review(): t = threading.Thread(target=_run_review, daemon=True, name="bg-review") t.start() + def _build_memory_write_metadata( + self, + *, + write_origin: Optional[str] = None, + execution_context: Optional[str] = None, + task_id: Optional[str] = None, + tool_call_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Build provenance metadata for external memory-provider mirrors.""" + metadata: Dict[str, Any] = { + "write_origin": write_origin or getattr(self, "_memory_write_origin", "assistant_tool"), + "execution_context": ( + execution_context + or getattr(self, "_memory_write_context", "foreground") + ), + "session_id": self.session_id or "", + "parent_session_id": self._parent_session_id or "", + "platform": self.platform or os.environ.get("HERMES_SESSION_SOURCE", "cli"), + "tool_name": "memory", + } + if task_id: + metadata["task_id"] = task_id + if tool_call_id: + metadata["tool_call_id"] = tool_call_id + return {k: v for k, v in metadata.items() if v not in (None, "")} + def _apply_persist_user_message_override(self, messages: List[Dict]) -> None: """Rewrite the current-turn user message before persistence/return. @@ -7748,6 +7779,19 @@ def flush_memories(self, messages: list = None, min_turns: int = None): old_text=args.get("old_text"), store=self._memory_store, ) + if self._memory_manager and args.get("action") in ("add", "replace"): + try: + self._memory_manager.on_memory_write( + args.get("action", ""), + flush_target, + args.get("content", ""), + metadata=self._build_memory_write_metadata( + write_origin="memory_flush", + execution_context="flush_memories", + ), + ) + except Exception: + pass if not self.quiet_mode: print(f" 🧠 Memory flush: saved to {args.get('target', 'memory')}") except Exception as e: @@ -7968,6 +8012,10 @@ def _invoke_tool(self, function_name: str, function_args: dict, effective_task_i function_args.get("action", ""), target, function_args.get("content", ""), + metadata=self._build_memory_write_metadata( + task_id=effective_task_id, + tool_call_id=tool_call_id, + ), ) except Exception: pass @@ -8479,6 +8527,10 @@ def _execute_tool_calls_sequential(self, assistant_message, messages: list, effe function_args.get("action", ""), target, function_args.get("content", ""), + metadata=self._build_memory_write_metadata( + task_id=effective_task_id, + tool_call_id=getattr(tool_call, "id", None), + ), ) except Exception: pass diff --git a/tests/agent/test_memory_provider.py b/tests/agent/test_memory_provider.py index 5cd0d8ab41367..ca39da70f0819 100644 --- a/tests/agent/test_memory_provider.py +++ b/tests/agent/test_memory_provider.py @@ -77,6 +77,13 @@ def on_memory_write(self, action, target, content): self.memory_writes.append((action, target, content)) +class MetadataMemoryProvider(FakeMemoryProvider): + """Provider that opts into write metadata.""" + + def on_memory_write(self, action, target, content, metadata=None): + self.memory_writes.append((action, target, content, metadata or {})) + + # --------------------------------------------------------------------------- # MemoryProvider ABC tests # --------------------------------------------------------------------------- @@ -862,6 +869,51 @@ def test_on_memory_write_add(self): mgr.on_memory_write("add", "memory", "new fact") assert p.memory_writes == [("add", "memory", "new fact")] + def test_on_memory_write_metadata_passed_to_opt_in_provider(self): + """Providers that accept metadata receive structured write provenance.""" + mgr = MemoryManager() + p = MetadataMemoryProvider("ext") + mgr.add_provider(p) + + mgr.on_memory_write( + "add", + "memory", + "new fact", + metadata={ + "write_origin": "assistant_tool", + "execution_context": "foreground", + "session_id": "sess-1", + }, + ) + + assert p.memory_writes == [ + ( + "add", + "memory", + "new fact", + { + "write_origin": "assistant_tool", + "execution_context": "foreground", + "session_id": "sess-1", + }, + ) + ] + + def test_on_memory_write_metadata_keeps_legacy_provider_compatible(self): + """Old 3-arg providers keep working when the manager receives metadata.""" + mgr = MemoryManager() + p = FakeMemoryProvider("ext") + mgr.add_provider(p) + + mgr.on_memory_write( + "add", + "user", + "legacy provider fact", + metadata={"write_origin": "assistant_tool"}, + ) + + assert p.memory_writes == [("add", "user", "legacy provider fact")] + def test_on_memory_write_replace(self): """on_memory_write fires for 'replace' actions.""" mgr = MemoryManager() diff --git a/tests/run_agent/test_flush_memories_codex.py b/tests/run_agent/test_flush_memories_codex.py index 28fbf550db3f9..03555d6ff25d3 100644 --- a/tests/run_agent/test_flush_memories_codex.py +++ b/tests/run_agent/test_flush_memories_codex.py @@ -209,6 +209,31 @@ def test_flush_executes_memory_tool_calls(self, monkeypatch): assert call_kwargs.kwargs["target"] == "notes" assert "dark mode" in call_kwargs.kwargs["content"] + def test_flush_bridges_memory_write_metadata(self, monkeypatch): + """Flush memory writes notify external providers with flush provenance.""" + agent = _make_agent(monkeypatch, api_mode="chat_completions", provider="openrouter") + agent._memory_manager = MagicMock() + agent.session_id = "sess-flush" + agent.platform = "cli" + + mock_response = _chat_response_with_memory_call() + + with patch("agent.auxiliary_client.call_llm", return_value=mock_response): + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + {"role": "user", "content": "Note this"}, + ] + with patch("tools.memory_tool.memory_tool", return_value="Saved."): + agent.flush_memories(messages) + + agent._memory_manager.on_memory_write.assert_called_once() + call_kwargs = agent._memory_manager.on_memory_write.call_args + assert call_kwargs.args[:3] == ("add", "notes", "User prefers dark mode.") + assert call_kwargs.kwargs["metadata"]["write_origin"] == "memory_flush" + assert call_kwargs.kwargs["metadata"]["execution_context"] == "flush_memories" + assert call_kwargs.kwargs["metadata"]["session_id"] == "sess-flush" + def test_flush_strips_artifacts_from_messages(self, monkeypatch): """After flush, the flush prompt and any response should be removed from messages.""" agent = _make_agent(monkeypatch, api_mode="chat_completions", provider="openrouter")