From 731e8c6eee624d8406d4ca0089e725585905484d Mon Sep 17 00:00:00 2001 From: Keswyk Cheung <41133441+zhangbububu@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:08:32 +0800 Subject: [PATCH] feat(memory): bulk-load sample for registered users (#4066) * feat(memory): bulk-load sample for registered users * fix(memory): harden all-users sample loader --------- Co-authored-by: Willem Jiang --- README.md | 22 +- backend/app/gateway/auth/repositories/base.py | 5 + .../app/gateway/auth/repositories/sqlite.py | 6 + backend/docs/MEMORY_SETTINGS_REVIEW.md | 39 +++- .../harness/deerflow/agents/memory/AGENTS.md | 3 + backend/tests/test_auth.py | 25 ++ backend/tests/test_load_memory_sample.py | 213 ++++++++++++++++++ scripts/load_memory_sample.py | 152 +++++++++++-- 8 files changed, 429 insertions(+), 36 deletions(-) create mode 100644 backend/tests/test_load_memory_sample.py diff --git a/README.md b/README.md index a172988ea3..73124751b1 100644 --- a/README.md +++ b/README.md @@ -472,19 +472,23 @@ such a checkout, use `bash ./scripts/.sh ...`. ``` Reads the configured sandbox image from UTF-8 `config.yaml`, with or without a leading BOM, using LF or CRLF line endings. -4. **(Optional) Load sample memory data for local review**: - ```bash - python scripts/load_memory_sample.py - ``` - This copies the sample fixture into the default local runtime memory file so reviewers can immediately test `Settings > Memory`. - See [backend/docs/MEMORY_SETTINGS_REVIEW.md](backend/docs/MEMORY_SETTINGS_REVIEW.md) for the shortest review flow. - -5. **Start services**: +4. **Start services**: ```bash make dev ``` -6. **Access**: http://localhost:2026 +5. **Access**: http://localhost:2026 + +6. **(Optional) Load sample memory data for local review**: open `Settings > Memory`, click **Import memory**, and select `backend/docs/memory-settings-sample.json`. The browser imports into the signed-in user's memory. + + To replace memory for every registered user in a disposable review environment: + + ```bash + cd backend + uv run python ../scripts/load_memory_sample.py --all-users + ``` + + Bulk mode supports SQLite/PostgreSQL user registries, creates timestamped backups under `.deer-flow/memory-sample-backups/`, and rejects the non-persistent `database.backend: memory` mode. See [backend/docs/MEMORY_SETTINGS_REVIEW.md](backend/docs/MEMORY_SETTINGS_REVIEW.md) for the complete review flow. Local services always use their internal ports (`8001`, `3000`, and `2026`). The root `.env` variable `PORT` configures only the published Docker ingress; diff --git a/backend/app/gateway/auth/repositories/base.py b/backend/app/gateway/auth/repositories/base.py index 3ec18f75cd..c34751a32b 100644 --- a/backend/app/gateway/auth/repositories/base.py +++ b/backend/app/gateway/auth/repositories/base.py @@ -88,6 +88,11 @@ class UserRepository(ABC): """Return total number of registered users.""" raise NotImplementedError + @abstractmethod + async def list_user_ids(self) -> list[str]: + """Return every registered user ID in deterministic creation order.""" + raise NotImplementedError + @abstractmethod async def count_admin_users(self) -> int: """Return number of users with system_role == 'admin'.""" diff --git a/backend/app/gateway/auth/repositories/sqlite.py b/backend/app/gateway/auth/repositories/sqlite.py index b722771c62..1de9f2195f 100644 --- a/backend/app/gateway/auth/repositories/sqlite.py +++ b/backend/app/gateway/auth/repositories/sqlite.py @@ -284,6 +284,12 @@ class SQLiteUserRepository(UserRepository): async with self._sf() as session: return await session.scalar(stmt) or 0 + async def list_user_ids(self) -> list[str]: + stmt = select(UserRow.id).order_by(UserRow.created_at, UserRow.id) + async with self._sf() as session: + result = await session.scalars(stmt) + return list(result) + async def count_admin_users(self) -> int: stmt = select(func.count()).select_from(UserRow).where(UserRow.system_role == "admin") async with self._sf() as session: diff --git a/backend/docs/MEMORY_SETTINGS_REVIEW.md b/backend/docs/MEMORY_SETTINGS_REVIEW.md index 31cb115346..31144ad447 100644 --- a/backend/docs/MEMORY_SETTINGS_REVIEW.md +++ b/backend/docs/MEMORY_SETTINGS_REVIEW.md @@ -20,18 +20,38 @@ Use this when reviewing the Memory Settings add/edit flow locally with the fewes If you already have DeerFlow running locally, you can reuse that existing setup. -2. Load the sample memory fixture. - - ```bash - python scripts/load_memory_sample.py - ``` - -3. Open `Settings > Memory`. +2. Open `Settings > Memory`. Default local URLs: - App: `http://localhost:2026` - Local frontend-only fallback: `http://localhost:3000` +3. Click **Import memory** and select: + + ```text + backend/docs/memory-settings-sample.json + ``` + + The browser imports the fixture into the currently signed-in user's memory. + +## Bulk Local Review + +To replace memory for every registered user in a disposable review environment: + +```bash +cd backend +uv run python ../scripts/load_memory_sample.py --all-users +``` + +This command: + +- supports SQLite and PostgreSQL registered-user databases; +- creates timestamped backups under `.deer-flow/memory-sample-backups/` before replacing memory; +- imports through the configured memory storage provider; and +- rejects `database.backend: memory`, which has no persistent user registry. + +Bulk loading replaces every registered user's memory. Do not run it against an environment containing memory you need to preserve. Pass `--no-backup` only when losing the existing memory is acceptable. + ## Minimal Manual Test 1. Click `Add fact`. @@ -58,6 +78,7 @@ Use this when reviewing the Memory Settings add/edit flow locally with the fewes ## Fixture Files - Sample fixture: `backend/docs/memory-settings-sample.json` -- Default local runtime target: `backend/.deer-flow/memory.json` +- Per-user file-storage target: `backend/.deer-flow/users/{user_id}/memory.json` +- Bulk backups: `backend/.deer-flow/memory-sample-backups/{timestamp}/{user_id}.json` -The loader script creates a timestamped backup automatically before overwriting an existing runtime memory file. +For an explicit one-file copy, pass `--target PATH` to the loader. There is no implicit target because authenticated sessions use different per-user paths. diff --git a/backend/packages/harness/deerflow/agents/memory/AGENTS.md b/backend/packages/harness/deerflow/agents/memory/AGENTS.md index c02e064367..4f1bd3409c 100644 --- a/backend/packages/harness/deerflow/agents/memory/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/memory/AGENTS.md @@ -37,6 +37,9 @@ Only `AuthMiddleware` can authorize the internal owner header. No-auth mode uses `DEFAULT_USER_ID`, which is `"default"`. An absolute `storage_path` opts out of the default per-user root. +`UserRepository.list_user_ids()` returns registered IDs in creation order for +administrative workflows such as `scripts/load_memory_sample.py --all-users`. +The bulk loader requires persistent database mode and backs up each user's memory by default. DeerMem uses this layout: diff --git a/backend/tests/test_auth.py b/backend/tests/test_auth.py index 0f58060944..4c348b7611 100644 --- a/backend/tests/test_auth.py +++ b/backend/tests/test_auth.py @@ -425,6 +425,31 @@ def test_sqlite_round_trip_new_fields(): asyncio.run(_run()) +def test_user_repository_lists_registered_user_ids(tmp_path): + import asyncio + + from app.gateway.auth.repositories.sqlite import SQLiteUserRepository + + async def _run() -> None: + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + + await init_engine( + "sqlite", + url=f"sqlite+aiosqlite:///{tmp_path}/scratch.db", + sqlite_dir=str(tmp_path), + ) + try: + repo = SQLiteUserRepository(get_session_factory()) + first = await repo.create_user(User(email="first@test.com", password_hash="hash")) + second = await repo.create_user(User(email="second@test.com", password_hash="hash")) + + assert await repo.list_user_ids() == [str(first.id), str(second.id)] + finally: + await close_engine() + + asyncio.run(_run()) + + # ── IntegrityError classification (OAuth conflict vs. everything else) ────── # # Regression coverage for a misclassification bug: create_user's diff --git a/backend/tests/test_load_memory_sample.py b/backend/tests/test_load_memory_sample.py new file mode 100644 index 0000000000..2ceb18a9ee --- /dev/null +++ b/backend/tests/test_load_memory_sample.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +SCRIPT_PATH = Path(__file__).resolve().parents[2] / "scripts" / "load_memory_sample.py" +SPEC = importlib.util.spec_from_file_location("load_memory_sample", SCRIPT_PATH) +assert SPEC is not None and SPEC.loader is not None +loader = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(loader) + + +def test_parse_args_requires_one_target_mode(tmp_path): + with pytest.raises(SystemExit): + loader.parse_args(tmp_path, []) + + +def test_parse_args_rejects_target_with_all_users(tmp_path): + with pytest.raises(SystemExit): + loader.parse_args(tmp_path, ["--target", "memory.json", "--all-users"]) + + +def test_parse_args_accepts_all_users(tmp_path): + args = loader.parse_args(tmp_path, ["--all-users"]) + + assert args.all_users is True + assert args.target is None + + +def test_load_sample_for_users_backs_up_every_user_before_import(tmp_path): + events = [] + current = { + "u1": {"facts": [{"id": "old-1"}]}, + "u2": {"facts": [{"id": "old-2"}]}, + } + + def load_memory(*, user_id): + events.append(("load", user_id)) + return current[user_id] + + def import_memory(sample, *, user_id): + assert (tmp_path / f"{user_id}.json").exists() + assert all((tmp_path / f"{uid}.json").exists() for uid in current) + events.append(("import", user_id)) + return sample + + count = loader.load_sample_for_users( + {"facts": [{"id": "sample"}]}, + ["u1", "u2"], + backup_root=tmp_path, + no_backup=False, + load_memory=load_memory, + import_memory=import_memory, + ) + + assert count == 2 + assert events == [ + ("load", "u1"), + ("load", "u2"), + ("import", "u1"), + ("import", "u2"), + ] + + +def test_load_sample_for_users_with_no_users_is_a_noop(tmp_path): + assert ( + loader.load_sample_for_users( + {"facts": []}, + [], + backup_root=tmp_path, + no_backup=False, + load_memory=lambda **_: pytest.fail("unexpected load"), + import_memory=lambda *_args, **_kwargs: pytest.fail("unexpected import"), + ) + == 0 + ) + + +def test_load_sample_for_users_stops_after_import_failure(tmp_path): + imported = [] + + def fail_second(_sample, *, user_id): + imported.append(user_id) + if user_id == "u2": + raise OSError("boom") + + with pytest.raises(OSError, match="boom"): + loader.load_sample_for_users( + {"facts": []}, + ["u1", "u2", "u3"], + backup_root=tmp_path, + no_backup=True, + load_memory=lambda **_: {}, + import_memory=fail_second, + ) + + assert imported == ["u1", "u2"] + + +def test_load_sample_for_users_rejects_unpersisted_sample(tmp_path): + with pytest.raises(OSError, match="not persisted for user u1"): + loader.load_sample_for_users( + {"facts": [{"id": "sample"}]}, + ["u1"], + backup_root=tmp_path, + no_backup=True, + load_memory=lambda **_: {"facts": []}, + import_memory=lambda *_args, **_kwargs: {"facts": []}, + ) + + +def test_load_sample_for_all_users_uses_runtime_config_resolution(monkeypatch, tmp_path): + import asyncio + + import deerflow.config.app_config as app_config + + config_arguments = [] + + def from_file(config_path=None): + config_arguments.append(config_path) + return SimpleNamespace(database=SimpleNamespace(backend="memory")) + + monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(tmp_path / "review-config.yaml")) + monkeypatch.setattr(app_config.AppConfig, "from_file", from_file) + + with pytest.raises(SystemExit, match="sqlite or postgres"): + asyncio.run( + loader.load_sample_for_all_users(SCRIPT_PATH.parents[1], {"facts": []}, no_backup=True), + ) + + assert config_arguments == [None] + + +def test_load_sample_for_all_users_uses_configured_memory_manager(monkeypatch, tmp_path): + import asyncio + + import app.gateway.auth.repositories.sqlite as sqlite_repository + import deerflow.agents.memory.manager as memory_manager + import deerflow.config.app_config as app_config + import deerflow.config.paths as config_paths + import deerflow.persistence.engine as persistence_engine + + manager_factory_calls = 0 + loaded = [] + imported = [] + + class FakeMemoryManager: + def get_memory(self, *, user_id): + loaded.append(user_id) + return {"facts": [{"id": f"old-{user_id}"}]} + + def import_memory(self, sample, *, user_id): + imported.append((user_id, sample)) + return sample + + manager = FakeMemoryManager() + + def get_memory_manager(): + nonlocal manager_factory_calls + manager_factory_calls += 1 + return manager + + class FakeUserRepository: + def __init__(self, _session_factory): + pass + + async def list_user_ids(self): + return ["u1", "u2"] + + async def init_engine_from_config(_database): + pass + + async def close_engine(): + pass + + monkeypatch.setattr(memory_manager, "get_memory_manager", get_memory_manager) + monkeypatch.setattr(sqlite_repository, "SQLiteUserRepository", FakeUserRepository) + monkeypatch.setattr( + app_config.AppConfig, + "from_file", + lambda *_args: SimpleNamespace(database=SimpleNamespace(backend="sqlite")), + ) + monkeypatch.setattr(config_paths, "get_paths", lambda: SimpleNamespace(base_dir=tmp_path)) + monkeypatch.setattr(persistence_engine, "init_engine_from_config", init_engine_from_config) + monkeypatch.setattr(persistence_engine, "get_session_factory", object) + monkeypatch.setattr(persistence_engine, "close_engine", close_engine) + + sample = {"facts": [{"id": "sample"}]} + count, backup_root = asyncio.run( + loader.load_sample_for_all_users(SCRIPT_PATH.parents[1], sample, no_backup=False), + ) + + assert count == 2 + assert manager_factory_calls == 1 + assert loaded == ["u1", "u2"] + assert imported == [("u1", sample), ("u2", sample)] + assert backup_root is not None + assert json.loads((backup_root / "u1.json").read_text())["facts"][0]["id"] == "old-u1" + assert json.loads((backup_root / "u2.json").read_text())["facts"][0]["id"] == "old-u2" + + +def test_require_persistent_database_rejects_memory(): + with pytest.raises(SystemExit, match="sqlite or postgres"): + loader.require_persistent_database("memory") + + +@pytest.mark.parametrize("backend", ["sqlite", "postgres"]) +def test_require_persistent_database_accepts_persistent_backends(backend): + loader.require_persistent_database(backend) diff --git a/scripts/load_memory_sample.py b/scripts/load_memory_sample.py index 568951e8f5..b5c4feb4fe 100644 --- a/scripts/load_memory_sample.py +++ b/scripts/load_memory_sample.py @@ -4,23 +4,24 @@ from __future__ import annotations import argparse +import asyncio +import copy import json import shutil +import sys +from collections.abc import Callable from datetime import datetime from pathlib import Path +from typing import Any def default_source(repo_root: Path) -> Path: return repo_root / "backend" / "docs" / "memory-settings-sample.json" -def default_target(repo_root: Path) -> Path: - return repo_root / "backend" / ".deer-flow" / "memory.json" - - -def parse_args(repo_root: Path) -> argparse.Namespace: +def parse_args(repo_root: Path, argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Copy the Memory Settings sample data into the local runtime memory file.", + description="Load Memory Settings sample data into DeerFlow runtime memory.", ) parser.add_argument( "--source", @@ -28,36 +29,151 @@ def parse_args(repo_root: Path) -> argparse.Namespace: default=default_source(repo_root), help="Path to the sample JSON file.", ) - parser.add_argument( + target_group = parser.add_mutually_exclusive_group(required=True) + target_group.add_argument( "--target", type=Path, - default=default_target(repo_root), - help="Path to the runtime memory.json file.", + help="Path to one runtime memory.json file.", + ) + target_group.add_argument( + "--all-users", + action="store_true", + help="Replace memory for every registered database user.", ) parser.add_argument( "--no-backup", action="store_true", - help="Overwrite the target without writing a backup copy first.", + help="Overwrite memory without writing backup copies first.", ) - return parser.parse_args() + return parser.parse_args(argv) -def validate_json_file(path: Path) -> None: +def load_json_file(path: Path) -> dict[str, Any]: with path.open(encoding="utf-8") as handle: - json.load(handle) + data = json.load(handle) + if not isinstance(data, dict): + raise SystemExit(f"Memory sample must be a JSON object: {path}") + return data -def main() -> int: +def require_persistent_database(backend: str) -> None: + if backend == "memory": + raise SystemExit("--all-users requires database.backend sqlite or postgres") + + +def load_sample_for_users( + sample: dict[str, Any], + user_ids: list[str], + *, + backup_root: Path | None, + no_backup: bool, + load_memory: Callable[..., dict[str, Any]], + import_memory: Callable[..., dict[str, Any]], +) -> int: + if not user_ids: + return 0 + + sample_facts = sample.get("facts", []) + expected_ids = {str(fact["id"]) for fact in sample_facts if isinstance(fact, dict) and fact.get("id")} + expected_contents = {str(fact["content"]) for fact in sample_facts if isinstance(fact, dict) and not fact.get("id") and fact.get("content")} + + if not no_backup: + if backup_root is None: + raise ValueError("backup_root is required when backups are enabled") + existing = {user_id: load_memory(user_id=user_id) for user_id in user_ids} + backup_root.mkdir(parents=True, exist_ok=True) + for user_id, memory in existing.items(): + (backup_root / f"{user_id}.json").write_text( + json.dumps(memory, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + + for user_id in user_ids: + try: + imported = import_memory(copy.deepcopy(sample), user_id=user_id) + except OSError as exc: + raise OSError(f"Failed to import memory for user {user_id}: {exc}") from exc + imported_facts = imported.get("facts", []) if isinstance(imported, dict) else [] + imported_ids = {str(fact["id"]) for fact in imported_facts if isinstance(fact, dict) and fact.get("id")} + imported_contents = {str(fact["content"]) for fact in imported_facts if isinstance(fact, dict) and fact.get("content")} + if not expected_ids <= imported_ids or not expected_contents <= imported_contents: + raise OSError(f"Imported memory was not persisted for user {user_id}") + + return len(user_ids) + + +async def load_sample_for_all_users( + repo_root: Path, + sample: dict[str, Any], + *, + no_backup: bool, +) -> tuple[int, Path | None]: + backend_dir = repo_root / "backend" + sys.path.insert(0, str(backend_dir)) + sys.path.insert(0, str(backend_dir / "packages" / "harness")) + + from app.gateway.auth.repositories.sqlite import SQLiteUserRepository + from deerflow.agents.memory.manager import get_memory_manager + from deerflow.config.app_config import AppConfig + from deerflow.config.paths import get_paths + from deerflow.persistence.engine import ( + close_engine, + get_session_factory, + init_engine_from_config, + ) + + config = await asyncio.to_thread(AppConfig.from_file) + require_persistent_database(config.database.backend) + manager = await asyncio.to_thread(get_memory_manager) + + await init_engine_from_config(config.database) + try: + session_factory = get_session_factory() + if session_factory is None: + raise SystemExit("Registered-user persistence is unavailable") + + user_ids = await SQLiteUserRepository(session_factory).list_user_ids() + backup_root = None + if user_ids and not no_backup: + backup_root = get_paths().base_dir / "memory-sample-backups" / datetime.now().strftime("%Y%m%d-%H%M%S") + + count = await asyncio.to_thread( + load_sample_for_users, + sample, + user_ids, + backup_root=backup_root, + no_backup=no_backup, + load_memory=manager.get_memory, + import_memory=manager.import_memory, + ) + return count, backup_root + finally: + await close_engine() + + +def main(argv: list[str] | None = None) -> int: repo_root = Path(__file__).resolve().parents[1] - args = parse_args(repo_root) + args = parse_args(repo_root, argv) source = args.source.resolve() - target = args.target.resolve() - if not source.exists(): raise SystemExit(f"Sample file not found: {source}") - validate_json_file(source) + sample = load_json_file(source) + + if args.all_users: + count, backup_root = asyncio.run( + load_sample_for_all_users(repo_root, sample, no_backup=args.no_backup), + ) + print(f"Loaded sample memory for {count} registered user(s).") + if backup_root is not None: + print(f"Backups created under: {backup_root}") + else: + print("No backups created.") + return 0 + + assert args.target is not None + target = args.target.resolve() target.parent.mkdir(parents=True, exist_ok=True) backup_path: Path | None = None