Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
171 changes: 171 additions & 0 deletions mempalace/mcp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
Tools (write):
mempalace_add_drawer — file verbatim content into a wing/room
mempalace_delete_drawer — remove a drawer by ID
mempalace_delete_by_source — bulk-remove all drawers mined from one source_file

Tools (maintenance):
mempalace_reconnect — force cache invalidation and reconnect after external writes
Expand Down Expand Up @@ -2060,6 +2061,158 @@ def _run():
_metadata_cache = None


def _purge_source_closets(source_file: str, *, commit: bool) -> int:
"""Count, and optionally delete, closets matching ``source_file`` exactly.

The closets collection is the searchable AAAK index layer; it is keyed by
``source_file`` independently of the drawers collection, so a drawer-only
delete would strand stale index pointers at the deleted source (#1722).
Mirrors the closet-purge step in :func:`mempalace.sync.sync_palace` and the
re-mine purge in :func:`mempalace.palace.purge_file_closets`.

Best-effort: a missing or unavailable closet collection yields 0 and never
raises, so it can never abort a drawer delete that has already committed.
Deletion is pushed down via ``delete(where=...)`` so it survives palaces
larger than the 10k ``get()`` truncation; the returned count is the (best
effort) number of matching closets observed before the delete.
"""
from .palace import get_closets_collection

try:
closets_col = get_closets_collection(_config.palace_path, create=False)
except Exception as exc:
logger.warning("Closet purge skipped (collection unavailable): %s", exc)
return 0
if closets_col is None:
return 0
try:
ids = closets_col.get(where={"source_file": source_file}, include=[]).get("ids") or []
count = len(ids)
if commit and count:
closets_col.delete(where={"source_file": source_file})
return count
except Exception as exc:
logger.warning("Closet purge failed for %s: %s", source_file, exc)
return 0


def tool_delete_by_source(source_file: str, dry_run: bool = True):
"""Delete every drawer whose ``source_file`` metadata matches exactly.

Bulk cleanup for the contamination case in #1722, where benchmark/eval
files (ShareGPT dumps, ``results_mempal_*.jsonl``, language config JSON)
get mined into the same wing as real user data and drown out semantic
search. Previously the only recourse was hand-rolled SQLite ``DELETE``
against ``chroma.sqlite3``.

Matching is exact on the stored ``source_file`` value and pushed down to
the backend via ``delete(where=...)`` — the same idiom used by the miner
and diary ingest paths — so there is no client-side id list and the
SQLite "too many variables" limit cannot be hit, regardless of how many
drawers share the source (the reporter had 55k).

Also purges the matching closets (the AAAK index layer) so deleting the
drawers doesn't strand stale index pointers at the dead source (#1722).

Defaults to a dry run: it reports the drawer match count, the closet match
count, and a small sample so the caller can confirm the blast radius before
anything is removed. Pass ``dry_run=False`` to commit the deletion
(irreversible).
"""
global _metadata_cache
if not isinstance(source_file, str) or not source_file.strip():
return {"success": False, "error": "source_file must be a non-empty string"}
# Mirror the ingestion-side normalization (tool_add_drawer strips lone
# surrogates from source_file before storing) so exact matching still hits
# rows mined from non-ASCII paths that arrived via a cp1252 stdin (#1488).
source_file = strip_lone_surrogates(source_file)

col = _get_collection()
if not col:
return _collection_error_or_no_palace()

where = {"source_file": source_file}
try:
# Paginated to survive palaces larger than the 10k get() truncation.
metas = _fetch_all_metadata(col, where=where)
except Exception as e:
return {"success": False, "error": str(e)}

match_count = len(metas)
# Distinct (wing, room) pairs so the caller sees where the hits live.
sample = []
seen = set()
for meta in metas:
meta = _safe_meta(meta)
# Default missing wing/room to "" for consistency with the rest of the
# file (drawers are always stored with both, but be defensive).
wing = meta.get("wing", "")
room = meta.get("room", "")
key = (wing, room)
if key in seen:
continue
seen.add(key)
sample.append({"wing": wing, "room": room})
if len(sample) >= 5:
break
Comment on lines +2145 to +2157

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If wing or room metadata is missing or None, meta.get() will return None. To maintain consistency with other tools in this file (which default missing wing and room to ""), we should default them to "" before adding them to the sample list.

Suggested change
for meta in metas:
meta = _safe_meta(meta)
key = (meta.get("wing"), meta.get("room"))
if key in seen:
continue
seen.add(key)
sample.append({"wing": meta.get("wing"), "room": meta.get("room")})
if len(sample) >= 5:
break
for meta in metas:
meta = _safe_meta(meta)
wing = meta.get("wing", "")
room = meta.get("room", "")
key = (wing, room)
if key in seen:
continue
seen.add(key)
sample.append({"wing": wing, "room": room})
if len(sample) >= 5:
break


if dry_run:
closet_match_count = _purge_source_closets(source_file, commit=False)
return {
"success": True,
"dry_run": True,
"source_file": source_file,
"match_count": match_count,
"closet_match_count": closet_match_count,
"sample": sample,
"hint": (
"No drawers were deleted. Re-run with dry_run=false to remove "
f"these {match_count} drawer(s) and {closet_match_count} index "
"entr(y/ies)."
if match_count
else "No drawers match this source_file."
),
}

if match_count == 0:
# Idempotent: deleting an absent source is a no-op, not an error.
return {
"success": True,
"dry_run": False,
"source_file": source_file,
"deleted": 0,
}

_wal_log(
"delete_by_source",
{"source_file": source_file, "match_count": match_count, "sample": sample},
)
try:
col.delete(where=where)
_metadata_cache = None
# Purge the matching closets too so the AAAK index doesn't keep stale
# pointers at the now-deleted drawers (#1722). Done after the drawer
# delete and intentionally best-effort: the drawers are already gone,
# so a closet-purge hiccup must not turn a successful delete into an
# error — it just leaves index cruft a later `repair` / re-mine clears.
closets_deleted = _purge_source_closets(source_file, commit=True)
logger.info(
"Deleted %d drawer(s) and %d closet(s) from source: %s",
match_count,
closets_deleted,
source_file,
)
return {
"success": True,
"dry_run": False,
"source_file": source_file,
"deleted": match_count,
"closets_deleted": closets_deleted,
}
except Exception as e:
return {"success": False, "error": str(e)}


def tool_sync(project_dir: str = None, wing: str = None, apply: bool = False):
"""Prune drawers whose source files are gitignored, missing, or moved (#1252)."""
global _metadata_cache
Expand Down Expand Up @@ -3173,6 +3326,24 @@ def tool_reconnect():
},
"handler": tool_mine,
},
"mempalace_delete_by_source": {
"description": "Bulk-delete every drawer mined from one source_file (exact match). Use to clean up benchmark/test data accidentally mined into a user wing (#1722). Returns a dry-run match count and sample by default; pass dry_run=false to commit. Irreversible.",
"input_schema": {
"type": "object",
"properties": {
"source_file": {
"type": "string",
"description": "Exact source_file metadata value to remove (e.g. the full path that was mined)",
},
"dry_run": {
"type": "boolean",
"description": "Preview the match count without deleting; default true. Pass false to actually delete.",
},
},
"required": ["source_file"],
},
"handler": tool_delete_by_source,
},
"mempalace_sync": {
"description": "Prune drawers whose source files are gitignored, deleted, or moved. Returns dry-run report by default; pass apply=true to commit deletions.",
"input_schema": {
Expand Down
188 changes: 188 additions & 0 deletions tests/test_mcp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -1929,6 +1929,194 @@ def test_update_drawer_chunked_logical_id_rewrites_group(monkeypatch, config, pa
assert listed["drawers"][0]["drawer_id"] == logical_id


# ── Delete by source (#1722) ────────────────────────────────────────────


class TestDeleteBySource:
"""``tool_delete_by_source`` — bulk cleanup of benchmark/test contamination (#1722)."""

def _seed(self, monkeypatch, config, palace_path, kg):
_patch_mcp_server(monkeypatch, config, kg)
_client, _col = _get_collection(palace_path, create=True)
del _client
from mempalace.mcp_server import tool_add_drawer

# Two drawers from a "benchmark" source, one from real user data.
tool_add_drawer(
wing="bench",
room="general",
content="ShareGPT yoga retreat conversation noise number one.",
source_file="results_mempal_hybrid_v4_session_1.jsonl",
)
tool_add_drawer(
wing="bench",
room="general",
content="ShareGPT coding job description noise number two.",
source_file="results_mempal_hybrid_v4_session_1.jsonl",
)
tool_add_drawer(
wing="clients",
room="webdesign",
content="GG Sauna Dachdecker real client memory that must survive.",
source_file="notes/clients.md",
)

def _seed_closets(self, palace_path):
"""Seed the AAAK index (closets) directly.

``tool_add_drawer`` never builds closets — those are a miner-side
artifact — so to exercise the closet purge we add them straight to the
collection, keyed by the same ``source_file`` the drawers use: two for
the benchmark source, one for the real-client source.
"""
from mempalace.palace import get_closets_collection

closets_col = get_closets_collection(palace_path, create=True)
closets_col.add(
ids=["bench_closet_01", "bench_closet_02", "client_closet_01"],
documents=[
"topic: yoga retreat | coding job",
"topic: more bench noise",
"topic: GG Sauna client",
],
metadatas=[
{"source_file": "results_mempal_hybrid_v4_session_1.jsonl"},
{"source_file": "results_mempal_hybrid_v4_session_1.jsonl"},
{"source_file": "notes/clients.md"},
],
)
return closets_col

def test_dry_run_reports_count_without_deleting(self, monkeypatch, config, palace_path, kg):
self._seed(monkeypatch, config, palace_path, kg)
from mempalace.mcp_server import tool_delete_by_source, tool_status

result = tool_delete_by_source("results_mempal_hybrid_v4_session_1.jsonl")
assert result["success"] is True
assert result["dry_run"] is True
assert result["match_count"] == 2
assert {"wing": "bench", "room": "general"} in result["sample"]
# Nothing removed — all three drawers still present.
assert tool_status()["total_drawers"] == 3

def test_dry_run_reports_closet_match_count(self, monkeypatch, config, palace_path, kg):
"""Dry run surfaces the closet blast radius (#1722) without deleting."""
self._seed(monkeypatch, config, palace_path, kg)
closets_col = self._seed_closets(palace_path)
from mempalace.mcp_server import tool_delete_by_source

result = tool_delete_by_source("results_mempal_hybrid_v4_session_1.jsonl")
assert result["dry_run"] is True
assert result["closet_match_count"] == 2
# Nothing removed — all three closets still present.
assert len(closets_col.get(include=[])["ids"]) == 3

def test_commit_deletes_only_matching_source(self, monkeypatch, config, palace_path, kg):
self._seed(monkeypatch, config, palace_path, kg)
from mempalace.mcp_server import tool_delete_by_source, tool_status

result = tool_delete_by_source("results_mempal_hybrid_v4_session_1.jsonl", dry_run=False)
assert result["success"] is True
assert result["dry_run"] is False
assert result["deleted"] == 2
# Only the real client drawer remains.
assert tool_status()["total_drawers"] == 1

def test_commit_purges_matching_closets(self, monkeypatch, config, palace_path, kg):
"""Deleting by source purges the matching closets too, so the AAAK
index keeps no stale pointers at the now-deleted drawers (#1722)."""
self._seed(monkeypatch, config, palace_path, kg)
closets_col = self._seed_closets(palace_path)
from mempalace.mcp_server import tool_delete_by_source

result = tool_delete_by_source("results_mempal_hybrid_v4_session_1.jsonl", dry_run=False)
assert result["success"] is True
assert result["deleted"] == 2
assert result["closets_deleted"] == 2
# The two benchmark closets are gone; the real-client closet survives.
remaining = closets_col.get(include=["metadatas"])
sources = {m["source_file"] for m in remaining["metadatas"]}
assert sources == {"notes/clients.md"}

def test_no_match_is_idempotent_not_error(self, monkeypatch, config, palace_path, kg):
self._seed(monkeypatch, config, palace_path, kg)
from mempalace.mcp_server import tool_delete_by_source, tool_status

result = tool_delete_by_source("does/not/exist.jsonl", dry_run=False)
assert result["success"] is True
assert result["deleted"] == 0
assert tool_status()["total_drawers"] == 3

def test_empty_source_file_rejected(self, monkeypatch, config, palace_path, kg):
self._seed(monkeypatch, config, palace_path, kg)
from mempalace.mcp_server import tool_delete_by_source

result = tool_delete_by_source(" ", dry_run=False)
assert result["success"] is False
assert "non-empty" in result["error"]

def test_non_string_source_rejected(self, monkeypatch, config, palace_path, kg):
"""A non-string source_file must return a clean error, not AttributeError."""
self._seed(monkeypatch, config, palace_path, kg)
from mempalace.mcp_server import tool_delete_by_source

result = tool_delete_by_source(123, dry_run=False)
assert result["success"] is False
assert "non-empty" in result["error"]

def test_matches_after_surrogate_normalization(self, monkeypatch, config, palace_path, kg):
"""source_file is stripped of lone surrogates on both ingest and delete,
so a path that arrived via a cp1252 stdin (#1488) still matches."""
_patch_mcp_server(monkeypatch, config, kg)
_client, _col = _get_collection(palace_path, create=True)
del _client
from mempalace.mcp_server import (
tool_add_drawer,
tool_delete_by_source,
tool_status,
)

# Lone low surrogate embedded in the path — add_drawer strips it.
raw_source = "noise\udce9_data.jsonl"
tool_add_drawer(
wing="bench",
room="general",
content="benchmark noise from a non-ASCII path",
source_file=raw_source,
)
assert tool_status()["total_drawers"] == 1

# Deleting with the same raw (un-stripped) string must still match.
result = tool_delete_by_source(raw_source, dry_run=False)
assert result["success"] is True
assert result["deleted"] == 1
assert tool_status()["total_drawers"] == 0

def test_registered_and_dispatchable(self, monkeypatch, config, palace_path, kg):
self._seed(monkeypatch, config, palace_path, kg)
from mempalace.mcp_server import handle_request

# Listed in tools/list
listed = handle_request({"method": "tools/list", "id": 1, "params": {}})
names = {t["name"] for t in listed["result"]["tools"]}
assert "mempalace_delete_by_source" in names

# Dispatches and defaults to dry-run (no destructive side effect)
resp = handle_request(
{
"method": "tools/call",
"id": 2,
"params": {
"name": "mempalace_delete_by_source",
"arguments": {"source_file": "results_mempal_hybrid_v4_session_1.jsonl"},
},
}
)
content = json.loads(resp["result"]["content"][0]["text"])
assert content["dry_run"] is True
assert content["match_count"] == 2


# ── KG Tools ────────────────────────────────────────────────────────────


Expand Down
Loading
Loading