agent-memory / tests /test_memory.py
edbeeching's picture
edbeeching HF Staff
Release portable agent memory client 0.1.0
d1d8733 verified
Raw History Blame Contribute Delete
17.2 kB
import importlib.util
import json
from pathlib import Path
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor
from types import SimpleNamespace
import pytest
ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location("memory", ROOT / "memory-use/scripts/memory.py")
m = importlib.util.module_from_spec(spec)
spec.loader.exec_module(m)
class Bucket:
"""An in-memory fake of bucket operations; never used by the production CLI."""
def __init__(self):
self.buckets = {}
self.calls = []
self.fail = None
self.fail_marker = False
self.unknown_put = False
def call(self, op, **kw):
self.calls.append((op, kw))
if self.fail:
raise m.MemoryError("Simulated " + self.fail, self.fail)
if op == "whoami":
return {"name": "owner", "orgs": ["team"]}
if op == "create":
if kw["bucket"] in self.buckets:
raise m.MemoryError("Already exists", "conflict")
self.buckets[kw["bucket"]] = {"private": True, "files": {}}
return {}
if op == "core":
return {"path": str(ROOT / "memory-use")}
if kw["bucket"] not in self.buckets:
raise m.MemoryError("Missing", "not-found")
bucket = self.buckets[kw["bucket"]]
files = bucket["files"]
if op == "info":
return {"id": kw["bucket"], "private": bucket["private"]}
if op == "put":
if self.fail_marker and kw["path"].endswith("/SKILL.md"):
raise m.MemoryError("Interrupted publication", "offline")
files[kw["path"]] = kw["text"]
if self.unknown_put:
raise m.MemoryError("Lost acknowledgment", "offline")
return {}
if op == "get":
if kw["path"] not in files:
raise m.MemoryError("Missing", "not-found")
return {"text": files[kw["path"]]}
if op == "list":
return [{"path": p, "size": len(t.encode())} for p, t in files.items() if p.startswith(kw["prefix"])]
raise AssertionError(op)
@pytest.fixture
def setup(tmp_path):
store = Bucket()
client = m.Memory(tmp_path / "state", store)
work = tmp_path / "project"
work.mkdir()
client.provision("project", work)
return client, store, work
def example_skill(tmp_path, body="Use a release checklist."):
folder = tmp_path / "release-check"
folder.mkdir(exist_ok=True)
(folder / "SKILL.md").write_text("---\nname: release-check\ndescription: Check releases.\n---\n\n" + body)
(folder / "references").mkdir(exist_ok=True)
(folder / "references/check.md").write_text("Verify the release artifact.")
return folder
def setup_runtime(client, tmp_path, store, work):
instructions = tmp_path / "AGENTS.md"
instructions.write_text("Keep existing conventions.\n")
return client.bootstrap("project", work, bucket="owner/agent-memory", skills_dir=tmp_path / "skills", instructions=instructions)
def test_two_machines_share_memories_and_keep_corrections(setup, tmp_path):
one, store, work = setup
saved = one.save("Use uv for this project.", "Package tooling", cwd=work)
two = m.Memory(tmp_path / "second-machine", store)
other_checkout = tmp_path / "other-checkout"
other_checkout.mkdir()
two.provision("project", other_checkout, "owner/agent-memory")
results = two.search("UV", cwd=other_checkout)
assert results["sync"] == "fresh" and len(results["results"]) == 1
assert "Use uv" in two.read(saved["id"], cwd=other_checkout)["text"]
correction = two.save("Correction: use pip. Supersedes " + saved["id"], "Tooling correction", cwd=other_checkout)
assert saved["id"] != correction["id"]
assert len(one.search(cwd=work)["results"]) == 2
def test_unknown_upload_retries_identical_bytes_after_configuration_change(setup):
client, store, work = setup
store.unknown_put = True
result = client.save("Verified result", "Checkpoint", kind="checkpoint", cwd=work)
assert result["status"] == "pending"
original = dict(store.buckets["owner/agent-memory"]["files"])
store.unknown_put = False
client.provision("other", work, create_bucket="owner/other")
assert client.retry(result["id"])["status"] == "saved"
assert store.buckets["owner/agent-memory"]["files"] == original
assert not store.buckets["owner/other"]["files"]
assert not client.pending()
def test_save_offline_is_durable_and_retryable(setup):
client, store, work = setup
store.fail = "offline"
result = client.save("Keep this finding", "Finding", cwd=work)
assert result["status"] == "pending"
assert client.pending()[0]["text"].endswith("Keep this finding\n")
restarted = m.Memory(client.home, store)
store.fail = None
assert restarted.retry(result["id"])["status"] == "saved"
def test_no_cleanup_before_remote_confirmation(setup, monkeypatch):
client, store, work = setup
original = m.write_json
def crash_on_receipt(path, value):
if Path(path).name == "last-upload.json":
raise OSError("Simulated disk failure after upload")
return original(path, value)
monkeypatch.setattr(m, "write_json", crash_on_receipt)
with pytest.raises(OSError):
client.save("Retain until acknowledged", "Crash case", cwd=work)
job = client.pending()[0]
assert len(store.buckets["owner/agent-memory"]["files"]) == 1
monkeypatch.setattr(m, "write_json", original)
client.retry(job["id"])
assert len(store.buckets["owner/agent-memory"]["files"]) == 1
assert not client.pending()
@pytest.mark.parametrize("failure", ["offline", "auth"])
def test_cached_search_distinguishes_failed_sync_and_does_not_prune(setup, failure):
client, store, work = setup
client.save("Cached evidence", "Evidence", cwd=work)
client.search(cwd=work)
store.fail = failure
result = client.search("evidence", cwd=work)
assert result["stale"] and len(result["results"]) == 1
assert result["sync"] == ("auth-error" if failure == "auth" else "offline")
def test_incomplete_listing_and_download_preserve_cache(setup, monkeypatch):
client, store, work = setup
old = client.save("Keep cached note", "Old", cwd=work)
client.search(cwd=work)
original = store.call
def interrupted(op, **kw):
if op == "list":
raise m.MemoryError("Pagination failed after page one", "offline")
return original(op, **kw)
monkeypatch.setattr(store, "call", interrupted)
assert len(client.search(cwd=work)["results"]) == 1
monkeypatch.setattr(store, "call", original)
store.buckets["owner/agent-memory"]["files"].pop(old["path"])
new = client.save("New note", "New", cwd=work)
def broken_download(op, **kw):
if op == "get" and kw["path"] == new["path"]:
raise m.MemoryError("Download interrupted", "offline")
return original(op, **kw)
monkeypatch.setattr(store, "call", broken_download)
assert client.search(cwd=work)["stale"]
assert (client.cache("owner/agent-memory") / old["path"]).exists()
assert not (client.cache("owner/agent-memory") / new["path"]).exists()
monkeypatch.setattr(store, "call", original)
assert not client.search(cwd=work)["stale"]
assert not (client.cache("owner/agent-memory") / old["path"]).exists()
def test_explicit_missing_bucket_never_creates_replacement(tmp_path):
store = Bucket()
client = m.Memory(tmp_path / "state", store)
with pytest.raises(m.MemoryError):
client.provision("project", tmp_path, "owner/missing")
assert not any(op == "create" for op, _ in store.calls)
def test_configured_missing_bucket_is_not_replaced(setup):
client, store, work = setup
store.buckets.clear()
store.calls.clear()
with pytest.raises(m.MemoryError):
client.provision("project", work)
assert not any(op == "create" for op, _ in store.calls)
def test_public_default_bucket_is_not_reconfigured(tmp_path):
store = Bucket()
store.buckets["owner/agent-memory"] = {"private": False, "files": {}}
with pytest.raises(m.MemoryError, match="private"):
m.Memory(tmp_path / "state", store).provision("project", tmp_path)
assert store.buckets["owner/agent-memory"]["private"] is False
assert not any(op == "create" for op, _ in store.calls)
def test_longest_project_ancestor_and_directory_boundaries(setup, tmp_path):
client, _, work = setup
nested = work / "nested"
nested.mkdir()
client.provision("nested", nested, "owner/agent-memory")
assert client.project(nested / "src")["project"] == "nested"
assert client.project(work / "other")["project"] == "project"
with pytest.raises(m.MemoryError, match="no memory project"):
client.save("Do not route incorrectly", "Unmapped", cwd=tmp_path / "project-other")
def test_bootstrap_repeat_preserves_instructions_and_verification(setup, tmp_path):
client, store, work = setup
first = setup_runtime(client, tmp_path, store, work)
instructions = tmp_path / "AGENTS.md"
instructions.write_text(instructions.read_text() + "\nMore unrelated instructions.\n")
second = client.bootstrap("project", work, bucket="owner/agent-memory", skills_dir=tmp_path / "skills", instructions=instructions)
assert first["verification_id"] == second["verification_id"]
assert len(store.buckets["owner/agent-memory"]["files"]) == 1
text = instructions.read_text()
assert text.startswith("Keep existing conventions.") and text.endswith("More unrelated instructions.\n")
assert text.count(m.BEGIN) == 1
def test_skill_revisions_partial_publication_and_frozen_retry(setup, tmp_path):
client, store, work = setup
folder = example_skill(tmp_path)
store.fail_marker = True
first = client.publish(folder, cwd=work)
assert first["status"] == "pending"
assert not client.search(skills=True, cwd=work)["results"]
original_text = (folder / "SKILL.md").read_text()
(folder / "SKILL.md").write_text(original_text + "\nA subsequent local edit.")
store.fail_marker = False
client.retry(first["id"])
assert store.buckets["owner/agent-memory"]["files"][first["path"] + "SKILL.md"] == original_text
second = client.publish(folder, cwd=work)
assert second["revision"] != first["revision"]
assert len(client.search(skills=True, cwd=work)["results"]) == 2
def test_install_preserves_modified_skill_and_explicit_replacement_backs_up(setup, tmp_path):
client, store, work = setup
setup_runtime(client, tmp_path, store, work)
folder = example_skill(tmp_path)
published = client.publish(folder, cwd=work)
client.install_shared("release-check", published["revision"], cwd=work)
installed = tmp_path / "skills/release-check/references/check.md"
installed.write_text("Local modifications")
with pytest.raises(m.MemoryError, match="Existing skill differs"):
client.install_shared("release-check", published["revision"], cwd=work)
assert installed.read_text() == "Local modifications"
result = client.install_shared("release-check", published["revision"], cwd=work, replace=True)
assert (Path(result["backup"]) / "references/check.md").read_text() == "Local modifications"
assert "release-check" in (tmp_path / "AGENTS.md").read_text()
@pytest.mark.parametrize("bad", ["../escape", "/absolute", "a/../../escape", "a\\b", "a//b", "./file", "C:/file"])
def test_unsafe_paths_rejected(bad):
with pytest.raises(m.MemoryError):
m.validate_skill({"SKILL.md": "---\nname: example\ndescription: Example\n---\nBody", bad: "bad"})
def test_symlink_publish_and_install_destination_rejected(setup, tmp_path):
client, store, work = setup
folder = example_skill(tmp_path)
(folder / "link").symlink_to(tmp_path / "outside")
with pytest.raises(m.MemoryError, match="symlink"):
client.publish(folder, cwd=work)
(folder / "link").unlink()
root = tmp_path / "linked-skills"
root.symlink_to(folder, target_is_directory=True)
with pytest.raises(m.MemoryError, match="Symlinks"):
client.install_files(m.skill_files(folder), "release-check", root, {})
def test_core_name_cannot_be_published_or_installed_from_bucket(setup):
client, _, work = setup
with pytest.raises(m.MemoryError, match="pinned Space"):
client.publish(ROOT / "memory-use", cwd=work)
with pytest.raises(m.MemoryError, match="cannot replace"):
client.install_shared("memory-use", "abc", cwd=work)
def test_reminder_has_no_network_or_transcript_reads(setup, tmp_path):
client, store, work = setup
store.calls.clear()
payload = {"hook_event_name": "UserPromptSubmit", "cwd": str(work), "transcript_path": "/must/not/read", "prompt": "Do not echo my message"}
result = m.hook_output(client, payload)
assert result["hookSpecificOutput"]["hookEventName"] == "UserPromptSubmit"
assert "Do not echo" not in json.dumps(result)
assert not store.calls
assert m.hook_output(client, {**payload, "cwd": str(tmp_path / "unmapped")}) == {}
def test_hook_merge_remove_preserves_unrelated_handlers(setup, tmp_path, monkeypatch):
client, store, work = setup
setup_runtime(client, tmp_path, store, work)
config = tmp_path / "hooks.json"
unrelated = {"type": "command", "command": "echo unrelated"}
initial = {"description": "Existing hooks", "hooks": {"SessionStart": [{"matcher": "resume", "hooks": [unrelated]}], "Stop": [{"hooks": [unrelated]}]}}
config.write_text(json.dumps(initial))
monkeypatch.setattr(m.subprocess, "run", lambda *a, **k: SimpleNamespace(stdout="codex-cli 0.162.1"))
m.hooks(client, "install", config)
m.hooks(client, "install", config)
value = json.loads(config.read_text())
assert len(value["hooks"]["SessionStart"]) == 2
m.hooks(client, "remove", config)
assert json.loads(config.read_text()) == initial
def test_credentials_rejected_and_errors_do_not_echo_sdk_secrets(setup, monkeypatch):
client, _, work = setup
token = "hf_" + "x" * 30
monkeypatch.setenv("HF_TOKEN", token)
with pytest.raises(m.MemoryError, match="credential"):
client.save(token, "Bad note", cwd=work)
exc = Exception("Authorization: " + token)
assert token not in str(m.error_from_exception(exc))
def test_cli_json_errors_and_hook_fail_open(tmp_path):
script = ROOT / "memory-use/scripts/memory.py"
command = [sys.executable, str(script), "--home", str(tmp_path / "state")]
result = subprocess.run(command + ["search", "query", "--json"], capture_output=True, text=True, cwd=tmp_path)
assert result.returncode == 1
assert json.loads(result.stdout)["code"] == "unmapped"
hook = subprocess.run(command + ["hook"], input="invalid JSON", capture_output=True, text=True)
assert hook.returncode == 0
assert "systemMessage" in json.loads(hook.stdout)
def test_native_network_child_has_hard_deadline(tmp_path, monkeypatch):
hanging = tmp_path / "hanging.py"
hanging.write_text("import time\ntime.sleep(10)\n")
monkeypatch.setattr(m, "__file__", str(hanging))
with pytest.raises(m.MemoryError, match="deadline"):
m.HFStore(timeout=0.05).call("whoami")
def test_concurrent_saves_survive_with_distinct_objects(setup):
client, store, work = setup
with ThreadPoolExecutor(max_workers=4) as pool:
futures = [pool.submit(client.save, f"Finding {i}", f"Concurrent {i}", cwd=work) for i in range(12)]
results = [f.result() for f in futures]
assert len({r["id"] for r in results}) == 12
assert all(r["status"] == "saved" for r in results)
assert len(store.buckets["owner/agent-memory"]["files"]) == 12
assert not client.pending()
def test_remote_traversal_skill_never_changes_installation(setup, tmp_path):
client, store, work = setup
setup_runtime(client, tmp_path, store, work)
prefix = "skills/hostile/revision/"
store.buckets["owner/agent-memory"]["files"].update({
prefix + "SKILL.md": "---\nname: hostile\ndescription: Test\n---\nInstructions",
prefix + "../../../escaped": "Bad path",
})
with pytest.raises(m.MemoryError, match="Unsafe"):
client.install_shared("hostile", "revision", cwd=work)
assert not (tmp_path / "skills/hostile").exists()
assert not (tmp_path / "escaped").exists()
def test_missing_authentication_has_clear_local_message(monkeypatch):
import huggingface_hub
monkeypatch.setattr(huggingface_hub, "get_token", lambda: None)
with pytest.raises(m.MemoryError, match="No HF credentials") as result:
m.hf_operation({"op": "whoami"})
assert result.value.code == "auth"
def test_fresh_read_never_silently_falls_back_to_cache(setup):
client, store, work = setup
saved = client.save("Cached", "Fresh check", cwd=work)
client.read(saved["id"], cwd=work)
store.fail = "offline"
assert client.read(saved["id"], cwd=work)["stale"]
with pytest.raises(m.MemoryError):
client.read(saved["id"], cwd=work, fresh=True)