diff --git a/gateway/run.py b/gateway/run.py index ba7ea43ad46ac..9caf4a7e77fea 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -9570,6 +9570,12 @@ async def _notify_long_running(): _previewed, ) response["already_sent"] = True + # Propagate final_response_sent separately so the gateway can + # distinguish "stream sent tool progress" from "stream sent the + # final answer". Fixes silent response drops when already_sent + # was set by earlier edits but the final text was never delivered. + if getattr(_sc, "final_response_sent", False): + response["final_response_sent"] = True return response diff --git a/gateway/stream_consumer.py b/gateway/stream_consumer.py index 5b529e63e8950..53cf9383da45e 100644 --- a/gateway/stream_consumer.py +++ b/gateway/stream_consumer.py @@ -326,7 +326,7 @@ async def run(self) -> None: self._last_sent_text = "" self._last_edit_time = time.monotonic() if got_done: - self._final_response_sent = self._already_sent + self._final_response_sent = True # chunks were just successfully sent return if got_segment_break: self._message_id = None @@ -372,12 +372,17 @@ async def run(self) -> None: if self._accumulated: if self._fallback_final_send: await self._send_fallback_final(self._accumulated) + self._final_response_sent = True elif current_update_visible: self._final_response_sent = True elif self._message_id: self._final_response_sent = await self._send_or_edit(self._accumulated) elif not self._already_sent: self._final_response_sent = await self._send_or_edit(self._accumulated) + elif self._last_sent_text.strip(): + # No remaining text but the user already saw the + # final answer via a previous progressive edit. + self._final_response_sent = True return if commentary_text is not None: @@ -414,13 +419,14 @@ async def run(self) -> None: except Exception: pass # Only confirm final delivery if the best-effort send above - # actually succeeded OR if the final response was already - # confirmed before we were cancelled. Previously this - # promoted any partial send (already_sent=True) to - # final_response_sent — which suppressed the gateway's - # fallback send even when only intermediate text (e.g. - # "Let me search…") had been delivered, not the real answer. - if _best_effort_ok and not self._final_response_sent: + # actually succeeded, or if there was accumulated content and + # an active message (meaning we attempted delivery before being + # cancelled). Previously this promoted any partial send + # (already_sent=True) to final_response_sent — which suppressed + # the gateway's fallback send even when only intermediate text + # (e.g. "Let me search…") had been delivered, not the real answer. + if (_best_effort_ok or (self._accumulated and self._message_id)) \ + and not self._final_response_sent: self._final_response_sent = True except Exception as e: logger.error("Stream consumer error: %s", e) diff --git a/run_agent.py b/run_agent.py index 325df9beb1d68..4cd60c014ffb7 100644 --- a/run_agent.py +++ b/run_agent.py @@ -1216,9 +1216,17 @@ def __init__( self._memory_flush_min_turns = int(mem_config.get("flush_min_turns", 6)) if self._memory_enabled or self._user_profile_enabled: from tools.memory_tool import MemoryStore + + # Build per-user namespace: "platform:user_id" + # Empty namespace = shared/global (CLI sessions, or when user_id is unknown) + _ns = "" + if getattr(self, '_user_id', None) and getattr(self, 'platform', None): + _ns = f"{self.platform}:{self._user_id}" + self._memory_store = MemoryStore( memory_char_limit=mem_config.get("memory_char_limit", 2200), user_char_limit=mem_config.get("user_char_limit", 1375), + namespace=_ns, ) self._memory_store.load_from_disk() except Exception: diff --git a/tests/tools/test_memory_tool.py b/tests/tools/test_memory_tool.py index 7f63aee1ebb0b..1eb458a4935f5 100644 --- a/tests/tools/test_memory_tool.py +++ b/tests/tools/test_memory_tool.py @@ -92,7 +92,7 @@ def test_system_override_blocked(self): @pytest.fixture() def store(tmp_path, monkeypatch): """Create a MemoryStore with temp storage.""" - monkeypatch.setattr("tools.memory_tool.get_memory_dir", lambda: tmp_path) + monkeypatch.setattr("tools.memory_tool.get_memory_dir", lambda ns="": tmp_path) s = MemoryStore(memory_char_limit=500, user_char_limit=300) s.load_from_disk() return s @@ -185,7 +185,7 @@ def test_remove_empty_old_text(self, store): class TestMemoryStorePersistence: def test_save_and_load_roundtrip(self, tmp_path, monkeypatch): - monkeypatch.setattr("tools.memory_tool.get_memory_dir", lambda: tmp_path) + monkeypatch.setattr("tools.memory_tool.get_memory_dir", lambda ns="": tmp_path) store1 = MemoryStore() store1.load_from_disk() @@ -198,7 +198,7 @@ def test_save_and_load_roundtrip(self, tmp_path, monkeypatch): assert "Alice, developer" in store2.user_entries def test_deduplication_on_load(self, tmp_path, monkeypatch): - monkeypatch.setattr("tools.memory_tool.get_memory_dir", lambda: tmp_path) + monkeypatch.setattr("tools.memory_tool.get_memory_dir", lambda ns="": tmp_path) # Write file with duplicates mem_file = tmp_path / "MEMORY.md" mem_file.write_text("duplicate entry\n§\nduplicate entry\n§\nunique entry") @@ -226,6 +226,122 @@ def test_empty_snapshot_returns_none(self, store): assert store.format_for_system_prompt("memory") is None +# ========================================================================= +# MemoryStore namespace isolation +# ========================================================================= + +class TestMemoryStoreNamespaceIsolation: + def test_namespace_isolates_users(self, tmp_path, monkeypatch): + """Two MemoryStore instances with different namespaces don't share entries.""" + def _get_dir(ns=""): + if ns: + d = tmp_path / ns.replace(":", "_") + else: + d = tmp_path + d.mkdir(parents=True, exist_ok=True) + return d + monkeypatch.setattr("tools.memory_tool.get_memory_dir", _get_dir) + + # User A's store + store_a = MemoryStore(memory_char_limit=500, user_char_limit=300, namespace="telegram:111") + store_a.load_from_disk() + store_a.add("memory", "User A secret") + + # User B's store — should NOT see User A's entry + store_b = MemoryStore(memory_char_limit=500, user_char_limit=300, namespace="telegram:222") + store_b.load_from_disk() + assert "User A secret" not in store_b.memory_entries + assert len(store_b.memory_entries) == 0 + + # Global store (no namespace) — should also NOT see User A's entry + store_global = MemoryStore(memory_char_limit=500, user_char_limit=300) + store_global.load_from_disk() + assert "User A secret" not in store_global.memory_entries + + # Verify User A's data is in the right directory + assert (tmp_path / "telegram_111" / "MEMORY.md").exists() + # User B hasn't written yet — write something to verify dir + store_b.add("memory", "User B secret") + assert (tmp_path / "telegram_222" / "MEMORY.md").exists() + # Confirm cross-contamination didn't happen + mem_a = (tmp_path / "telegram_111" / "MEMORY.md").read_text() + mem_b = (tmp_path / "telegram_222" / "MEMORY.md").read_text() + assert "User A secret" in mem_a + assert "User B secret" in mem_b + assert "User A secret" not in mem_b + assert "User B secret" not in mem_a + + +class TestMemoryMigrationFromSharedRoot: + """Migration: shared root MEMORY.md/USER.md → per-user namespace dir.""" + + def test_migration_copies_shared_files_to_new_namespace(self, tmp_path, monkeypatch): + """When a namespaced user loads memory for the first time and the + shared root still has files, they should be auto-copied.""" + root = tmp_path # shared root + user_dir = tmp_path / "telegram_5137755622" + (root / "MEMORY.md").write_text("Shared memory entry\n§\nShared entry 2") + (root / "USER.md").write_text("Shared user profile") + + def _get_dir(namespace=""): + if namespace: + return user_dir + return root + + monkeypatch.setattr("tools.memory_tool.get_memory_dir", _get_dir) + + store = MemoryStore(memory_char_limit=500, user_char_limit=300, namespace="telegram:5137755622") + store.load_from_disk() + + assert (user_dir / "MEMORY.md").exists() + assert (user_dir / "USER.md").exists() + assert "Shared memory entry" in (user_dir / "MEMORY.md").read_text() + assert "Shared user profile" in (user_dir / "USER.md").read_text() + # Shared root files should still exist + assert (root / "MEMORY.md").exists() + assert (root / "USER.md").exists() + # Entries should be loaded + assert "Shared memory entry" in store.memory_entries + assert "Shared user profile" in store.user_entries + + def test_migration_does_not_overwrite_existing_user_files(self, tmp_path, monkeypatch): + """If the user dir already has files, migration should not overwrite.""" + root = tmp_path + user_dir = tmp_path / "telegram_5137755622" + user_dir.mkdir() + (root / "MEMORY.md").write_text("Old shared memory") + (user_dir / "MEMORY.md").write_text("User's own memory") + + def _get_dir(namespace=""): + if namespace: + return user_dir + return root + + monkeypatch.setattr("tools.memory_tool.get_memory_dir", _get_dir) + + store = MemoryStore(memory_char_limit=500, user_char_limit=300, namespace="telegram:5137755622") + store.load_from_disk() + + assert "User's own memory" in (user_dir / "MEMORY.md").read_text() + assert "Old shared memory" not in (user_dir / "MEMORY.md").read_text() + + def test_migration_does_nothing_for_empty_namespace(self, tmp_path, monkeypatch): + """Global/CLI sessions (no namespace) should not trigger migration.""" + root = tmp_path + (root / "MEMORY.md").write_text("Shared memory") + + def _get_dir(namespace=""): + return root + + monkeypatch.setattr("tools.memory_tool.get_memory_dir", _get_dir) + + store = MemoryStore(memory_char_limit=500, user_char_limit=300) # no namespace + store.load_from_disk() + + # Should just load normally, no migration + assert "Shared memory" in store.memory_entries + + # ========================================================================= # memory_tool() dispatcher # ========================================================================= diff --git a/tools/memory_tool.py b/tools/memory_tool.py index eef64e709669b..b3792c4a4bd5c 100644 --- a/tools/memory_tool.py +++ b/tools/memory_tool.py @@ -50,9 +50,45 @@ # (HERMES_HOME env var changes) are always respected. The old module-level # constant was cached at import time and could go stale if a profile switch # happened after the first import. -def get_memory_dir() -> Path: - """Return the profile-scoped memories directory.""" - return get_hermes_home() / "memories" +def get_memory_dir(namespace: str = "") -> Path: + """Return the profile-scoped memories directory, optionally per-user.""" + base = get_hermes_home() / "memories" + if namespace: + # Sanitize namespace to be filesystem-safe + safe_ns = re.sub(r'[^\w\-_@\.\+]', '_', namespace) + return base / safe_ns + return base + + +# --------------------------------------------------------------------------- +# Migration: when a namespaced user dir is created for the first time and +# the shared root still has MEMORY.md / USER.md, automatically copy them +# into the user dir so no memories are lost after enabling per-user +# isolation. +# --------------------------------------------------------------------------- + +def _migrate_from_shared(namespace: str) -> None: + """Copy shared root memories into a newly-created namespaced directory. + + Only copies files that do NOT already exist in the target. The shared + root files are left untouched so CLI / other users can still read them. + """ + if not namespace: + return + mem_dir = get_memory_dir(namespace) + root_dir = get_memory_dir() # shared root + + for name in ("MEMORY.md", "USER.md"): + src = root_dir / name + dst = mem_dir / name + if dst.exists() or not src.exists(): + continue + try: + content = src.read_text(encoding="utf-8") + dst.write_text(content, encoding="utf-8") + logger.info("Migrated shared %s → %s", src, dst) + except (OSError, IOError) as exc: + logger.warning("Could not migrate %s: %s", src, exc) ENTRY_DELIMITER = "\n§\n" @@ -111,21 +147,28 @@ class MemoryStore: Never mutated mid-session. Keeps prefix cache stable. - memory_entries / user_entries: live state, mutated by tool calls, persisted to disk. Tool responses always reflect this live state. + + Per-user isolation: when a namespace (e.g. "platform:user_id") is provided, + memory files are stored under a subdirectory, preventing cross-user leakage. """ - def __init__(self, memory_char_limit: int = 2200, user_char_limit: int = 1375): + def __init__(self, memory_char_limit: int = 2200, user_char_limit: int = 1375, namespace: str = ""): self.memory_entries: List[str] = [] self.user_entries: List[str] = [] self.memory_char_limit = memory_char_limit self.user_char_limit = user_char_limit + self.namespace = namespace # e.g. "telegram:123456" — empty = shared/global # Frozen snapshot for system prompt -- set once at load_from_disk() self._system_prompt_snapshot: Dict[str, str] = {"memory": "", "user": ""} def load_from_disk(self): """Load entries from MEMORY.md and USER.md, capture system prompt snapshot.""" - mem_dir = get_memory_dir() + mem_dir = get_memory_dir(self.namespace) mem_dir.mkdir(parents=True, exist_ok=True) + # Migrate shared root memories on first load for a new namespaced user. + _migrate_from_shared(self.namespace) + self.memory_entries = self._read_file(mem_dir / "MEMORY.md") self.user_entries = self._read_file(mem_dir / "USER.md") @@ -176,9 +219,8 @@ def _file_lock(path: Path): pass fd.close() - @staticmethod - def _path_for(target: str) -> Path: - mem_dir = get_memory_dir() + def _path_for(self, target: str) -> Path: + mem_dir = get_memory_dir(self.namespace) if target == "user": return mem_dir / "USER.md" return mem_dir / "MEMORY.md" @@ -194,7 +236,8 @@ def _reload_target(self, target: str): def save_to_disk(self, target: str): """Persist entries to the appropriate file. Called after every mutation.""" - get_memory_dir().mkdir(parents=True, exist_ok=True) + mem_dir = get_memory_dir(self.namespace) + mem_dir.mkdir(parents=True, exist_ok=True) self._write_file(self._path_for(target), self._entries_for(target)) def _entries_for(self, target: str) -> List[str]: