diff --git a/CHANGELOG.md b/CHANGELOG.md index 74797966b4..3138c9c762 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,14 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), --- +## [3.3.5] — unreleased + +### Bug Fixes + +- **`mempalace_diary_read` silently dropped entries on agent-name case mismatch.** `tool_diary_write` stored the `agent` metadata verbatim after `sanitize_name`, which preserves case, while `tool_diary_read` filtered by exact match. Writing as `"Claude"` and reading as `"claude"` (or vice-versa) returned zero rows. Both endpoints now lowercase `agent_name` immediately after sanitization, so reads are case-insensitive and the default per-agent wing slug is stable across casings. **Behavior change:** entries written prior to this fix under mixed-case agent names will not match the new lowercase filter; run `mempalace repair` if you need to migrate legacy diary metadata. (#1243) + +--- + ## [3.3.4] — 2026-04-30 ### Added @@ -19,6 +27,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Bug Fixes +- **MCP server `tool_diary_write` SIGSEGV when EF default differs.** `mcp_server._get_collection` bypassed `ChromaBackend.get_collection` and called `client.get_collection` / `client.create_collection` without `embedding_function=`. ChromaDB 1.x does not persist the EF identity, so the MCP server silently bound chromadb's `DefaultEmbeddingFunction` while the miner bound `mempalace.embedding.get_embedding_function()`. On bleeding-edge interpreters (python 3.14 + chromadb 1.5.x) this SIGSEGV'd the MCP server on first `col.add()`. `_get_collection` now resolves and passes the EF explicitly. (#1299, follow-up to #1262/#1289) - **Cross-wing topic tunnels for hyphenated dir names.** `mempalace init` recorded the `topics_by_wing` registry key under the raw directory name (e.g. `mempalace-public`), while `mempalace.yaml`'s `wing` field used the lower-cased + separator-collapsed slug (`mempalace_public`). At mine time the miner read the slug from the yaml and missed the registry, so `_compute_topic_tunnels_for_wing` returned `0` silently. Real-world: any project whose folder contained a hyphen or space lost every topic tunnel. Producer side: `cmd_init`, `room_detector_local`, `miner.load_config` no-yaml fallback, and `convo_miner` now all route through a shared `normalize_wing_name()` in `config.py` so future writes use the same key. Lookup side: `palace_graph.create_tunnel`, `list_tunnels`, `follow_tunnels`, and `find_tunnels` normalize incoming wing names too, so existing palaces with raw-name keys on disk also recover. (#1194, #1195, #1197, follow-up to #1180) - **HNSW index bloat from repeated resize+persist cycles.** ChromaDB's HNSW segment was growing into the tens of GB on palaces past ~15K drawers because `link_lists.bin` was being re-allocated on every flush. Setting `hnsw:batch_size` and `hnsw:sync_threshold` on collection metadata via the new `_HNSW_BLOAT_GUARD` constant pins the segment to one allocation per batch instead. Empirical: a fresh 39,792-drawer palace went from 30 GB on disk and segfaulting `mempalace status` to 376 MB and instant. Migration note — already-bloated palaces still need a `mempalace repair` or full re-mine; HNSW config is honoured at collection-create time only. (#1191, supersedes #346) - **`max_seq_id` poisoning from old `_fix_blob_seq_ids` shim.** The 0.6.x → 1.5.x BLOB-to-INTEGER migration was running `int.from_bytes(blob, 'big')` over chromadb 1.5.x's native `b'\x11\x11' + ASCII-digit` `max_seq_id` format, yielding ~1.23e18 integers that silently suppressed every subsequent `embeddings_queue` write for the affected segment. The shim is now narrowed to the `embeddings` table only, with an additional defense-in-depth guard that skips sysdb-10-prefixed BLOBs even there. New `mempalace repair --mode max-seq-id` un-poisons existing palaces either from a pre-corruption sidecar DB (exact restore) or heuristically (`MAX(embeddings.seq_id)` over the owning collection). (#1135) diff --git a/hooks/README.md b/hooks/README.md index 469d231992..adec50b689 100644 --- a/hooks/README.md +++ b/hooks/README.md @@ -114,17 +114,23 @@ The `stop_hook_active` flag prevents infinite loops: block once → AI saves → ``` Context window getting full → Claude Code fires PreCompact ↓ - Find transcript (from input or session_id lookup) - ↓ - Auto-mine transcript → palace (tool output captured) - ↓ - {"decision": "block", "reason": "save tool output verbatim..."} + Hook may block (configurable) ↓ AI saves everything ↓ Compaction proceeds ``` +#### PreCompact behavior modes + +`mempal_precompact_hook.sh` supports `MEMPAL_PRECOMPACT_MODE`: + +- `block_once` — block once per session (default for non-Claude harnesses) +- `block` — always block +- `proceed` — never block (default when the harness is detected as Claude/Claude Code) + +The Claude default is `proceed` because some Claude setups don’t automatically retry compaction after a block, which can prevent compaction from happening at all. + No counting needed — compaction always warrants a save. The auto-mine captures raw tool output before the AI gets a chance to summarize it away. ## Debugging @@ -155,27 +161,4 @@ export MEMPAL_PYTHON="/usr/bin/python3" # system Python is fin export MEMPAL_PYTHON="$HOME/.venvs/mempalace/bin/python" # or your venv ``` -Resolution priority: `$MEMPAL_PYTHON` (if set and executable) → `$(command -v python3)` → bare `python3`. The interpreter only needs `json` and `sys` from the standard library — `mempalace` itself does not need to be installed in it. - -Note: the `mempalace mine` auto-ingest runs via the `mempalace` CLI, so that command also needs to be on the hook's `PATH`. Installing with `pipx install mempalace` or `uv tool install mempalace` puts it on a stable global location; otherwise extend the hook environment's `PATH` to include your venv's `bin/`. - -## Backfill Past Conversations - -The hooks only capture conversations going forward. To mine **past** Claude Code sessions into your palace, run a one-time backfill: - -```bash -mempalace mine ~/.claude/projects/ --mode convos -``` - -This scans all JSONL transcripts from previous sessions and files them into the `conversations` wing. On a typical developer machine with months of history, this can yield 50K–200K drawers. - -For Codex CLI sessions: -```bash -mempalace mine ~/.codex/sessions/ --mode convos -``` - -This only needs to be done once — after that, the hooks auto-mine each session as you go. - -## Cost - -**Zero extra tokens.** The hooks notify the AI that saves happened in the background — the AI doesn't need to write anything in the chat. All filing is handled automatically. Previous versions asked the AI to write diary entries and drawer content in the chat window, which cost ~$1/session in retransmitted tokens. +Resolution priority: `$MEMPAL_PYTHON` (if set and executable) → `python3` from the hook's `PATH`. diff --git a/hooks/mempal_enrich_knowledge.sh b/hooks/mempal_enrich_knowledge.sh new file mode 100755 index 0000000000..c6c3c39ae0 --- /dev/null +++ b/hooks/mempal_enrich_knowledge.sh @@ -0,0 +1,82 @@ +#!/bin/bash +# MEMPALACE KNOWLEDGE ENRICHMENT HELPER +# Add structured facts to MemPalace knowledge graph for better querying +# +# Usage: ./mempal_enrich_knowledge.sh "subject" "predicate" "object" [--valid-from DATE] [--source-file FILE] +# +# Examples: +# ./mempal_enrich_knowledge.sh "ananas-ai" "is_project_of" "ananas-platform" --valid-from "2026-01-01" +# ./mempal_enrich_knowledge.sh "ai-marketing" "uses_mcp" "algolia" --source-file "~/projects/ai-marketing/algolia.ts" + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MEMPALACE_SRC="$(dirname "$SCRIPT_DIR")" + +# Check minimum arguments +if [[ $# -lt 3 ]]; then + echo "Error: Subject, predicate, and object are required" + echo "Usage: $0 \"subject\" \"predicate\" \"object\" [--valid-from DATE] [--source-file FILE]" + exit 1 +fi + +SUBJECT="$1" +PREDICATE="$2" +OBJECT="$3" +shift 3 + +# Parse optional arguments +VALID_FROM="" +SOURCE_FILE="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --valid-from) + VALID_FROM="$2" + shift 2 + ;; + --source-file) + SOURCE_FILE="$2" + shift 2 + ;; + *) + echo "Unknown option: $1" + exit 1 + ;; + esac +done + +# Create temporary Python script +TMP_PY=$(mktemp) +cat > "$TMP_PY" << EOF +import sys +sys.path.insert(0, '$MEMPALACE_SRC') +from mempalace.knowledge_graph import KnowledgeGraph +import json +import os + +subject = '$SUBJECT' +predicate = '$PREDICATE' +object = '$OBJECT' +valid_from = '$VALID_FROM' if '$VALID_FROM' else None +source_file = '$SOURCE_FILE' if '$SOURCE_FILE' else None + +# Expand ~ in source_file +if source_file and source_file.startswith('~/'): + source_file = os.path.expanduser(source_file) + +try: + kg = KnowledgeGraph() + triple_id = kg.add_triple(subject, predicate, object, + valid_from=valid_from, + source_file=source_file) + print(json.dumps({'success': True, 'triple_id': triple_id}, indent=2)) +except Exception as e: + print(json.dumps({'success': False, 'error': str(e)}, indent=2)) + sys.exit(1) +EOF + +python3 "$TMP_PY" +EXIT_CODE=$? +rm -f "$TMP_PY" +exit $EXIT_CODE \ No newline at end of file diff --git a/hooks/mempal_opencode_hook.sh b/hooks/mempal_opencode_hook.sh new file mode 100755 index 0000000000..8d06480b19 --- /dev/null +++ b/hooks/mempal_opencode_hook.sh @@ -0,0 +1,100 @@ +#!/bin/bash +# MEMPALACE HOOK FOR OPENCODE +# Adapted from Claude Code hooks for use with OpenCode AI +# +# Install: Add to OpenCode's MCP config or create a custom command +# For OpenCode MCP integration, add to opencode.json: +# { +# "mcp": { +# "mempalace": { +# "type": "local", +# "command": "/home/zapostolski/.mempalace/src/hooks/mempal_opencode_hook.sh", +# "args": ["stop"] +# } +# } +# } +# +# Or run manually: echo '{}' | ./mempal_opencode_hook.sh stop + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_DIR="$(dirname "$SCRIPT_DIR")" +STATE_DIR="$HOME/.mempalace/hook_state" +SAVE_INTERVAL=10 +SESSION_ID="opencode-$(date +%Y%m%d-%H%M%S)" + +mkdir -p "$STATE_DIR" + +log() { + echo "[$(date '+%H:%M:%S')] $1" >> "$STATE_DIR/hook.log" +} + +run_save() { + log "OPENCODE SAVE TRIGGERED for session $SESSION_ID" + + # Output block reason to trigger save in OpenCode + cat << 'HOOKJSON' +{ + "decision": "block", + "reason": "AUTO-SAVE checkpoint. Save all new decisions, findings, milestones, blockers, architecture changes, code changes, and next steps from this session to your memory system. Be thorough - organize into wings/rooms. After saving, continue the conversation." +} +HOOKJSON +} + +run_precompact() { + log "OPENCODE PRE-COMPACT triggered" + + cat << 'HOOKJSON' +{ + "decision": "block", + "reason": "COMPACTION IMMINENT. Save ALL topics, decisions, quotes, code, and context from this session to memory. After compaction detailed context will be lost. Save everything, then allow compaction." +} +HOOKJSON +} + +handle_stop() { + local since_last="${1:-0}" + + if [ "$since_last" -ge "$SAVE_INTERVAL" ]; then + run_save + else + log "OpenCode: $since_last exchanges since last save (threshold: $SAVE_INTERVAL)" + echo '{}' + fi +} + +increment_counter() { + local counter_file="$STATE_DIR/opencode_counter" + local count=0 + if [ -f "$counter_file" ]; then + count=$(cat "$counter_file" 2>/dev/null || echo 0) + fi + count=$((count + 1)) + echo "$count" > "$counter_file" + echo "$count" +} + +main() { + local hook_type="${1:-stop}" + INPUT=$(cat) + + log "OpenCode hook triggered: $hook_type" + + case "$hook_type" in + stop) + local count + count=$(increment_counter) + handle_stop "$count" + ;; + precompact) + run_precompact + ;; + *) + log "Unknown hook type: $hook_type" + echo '{}' + ;; + esac +} + +main "$@" diff --git a/hooks/mempal_opencode_simple.sh b/hooks/mempal_opencode_simple.sh new file mode 100755 index 0000000000..1e7810dad9 --- /dev/null +++ b/hooks/mempal_opencode_simple.sh @@ -0,0 +1,77 @@ +#!/bin/bash +# MEMPALACE OPENCODE HOOK (No Python dependencies required) +# This hook logs checkpoints and can be run alongside mempalace's Codex/Claude hooks +# +# Usage: +# 1. Run manually after important sessions: ./mempal_opencode_simple.sh save +# 2. Or integrate via OpenCode custom commands +# +# This creates checkpoint files that can be mined later by mempalace + +set -euo pipefail + +STATE_DIR="$HOME/.mempalace/hook_state" +CHECKPOINT_DIR="$STATE_DIR/opencode_checkpoints" +SESSION_ID="opencode-$(date +%Y%m%d-%H%M%S)" +TIMESTAMP=$(date '+%Y-%m-%d %H:%M:%S') + +mkdir -p "$STATE_DIR" "$CHECKPOINT_DIR" + +log() { + echo "[$(date '+%H:%M:%S')] OPENCODE: $1" >> "$STATE_DIR/hook.log" +} + +create_checkpoint() { + local checkpoint_file="$CHECKPOINT_DIR/${SESSION_ID}.md" + local notes="${1:-}" + + log "Creating checkpoint: $SESSION_ID" + + cat > "$checkpoint_file" << EOF +# OpenCode Session Checkpoint + +**Session:** $SESSION_ID +**Timestamp:** $TIMESTAMP +**User:** Zharko Apostolski + +## Session Notes +$notes + +## Key Topics Discussed + +## Decisions Made + +## Code Changes + +## Next Steps + +--- +*Created by mempalace OpenCode hook* +EOF + + echo "Checkpoint saved: $checkpoint_file" + log "Checkpoint created: $checkpoint_file" +} + +main() { + local command="${1:-save}" + + case "$command" in + save) + log "Save command received" + create_checkpoint "Manual checkpoint from OpenCode session" + ;; + status) + log "Status check" + echo "OpenCode mempalace hook status:" + echo " Checkpoints: $(ls -1 "$CHECKPOINT_DIR" 2>/dev/null | wc -l)" + echo " Last: $(ls -t "$CHECKPOINT_DIR" 2>/dev/null | head -1)" + ;; + *) + echo "Usage: $0 {save|status}" + exit 1 + ;; + esac +} + +main "$@" diff --git a/hooks/mempal_precompact_hook.sh b/hooks/mempal_precompact_hook.sh index e811d65600..fc9c182680 100755 --- a/hooks/mempal_precompact_hook.sh +++ b/hooks/mempal_precompact_hook.sh @@ -4,13 +4,6 @@ # Claude Code "PreCompact" hook. Fires RIGHT BEFORE the conversation # gets compressed to free up context window space. # -# This is the safety net. When compaction happens, the AI loses detailed -# context about what was discussed. This hook forces one final save of -# EVERYTHING before that happens. -# -# Unlike the save hook (which triggers every N exchanges), this ALWAYS -# blocks — because compaction is always worth saving before. -# # === INSTALL === # Add to .claude/settings.local.json: # @@ -23,101 +16,26 @@ # }] # }] # } -# -# For Codex CLI, add to .codex/hooks.json: -# -# "PreCompact": [{ -# "type": "command", -# "command": "/absolute/path/to/mempal_precompact_hook.sh", -# "timeout": 30 -# }] -# -# === HOW IT WORKS === -# -# Claude Code sends JSON on stdin with: -# session_id — unique session identifier -# -# We always return decision: "block" with a reason telling the AI -# to save everything. After the AI saves, compaction proceeds normally. -# -# === MEMPALACE CLI === -# The hook ALWAYS mines the active conversation transcript synchronously -# before compaction (via `mempalace mine --mode convos`). -# MEMPAL_DIR is an *additional*, optional target for project files — it -# does not replace the conversation mine. -STATE_DIR="$HOME/.mempalace/hook_state" -mkdir -p "$STATE_DIR" +set -euo pipefail -# Optional: project directory (code / notes / docs) to also mine before -# compaction. Mined with `--mode projects`. The conversation transcript -# is always mined regardless — this is purely additive. -# Example: MEMPAL_DIR="$HOME/projects/my_app" -MEMPAL_DIR="" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MEMPALACE_SRC="$(dirname "$SCRIPT_DIR")" +export PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$MEMPALACE_SRC" -# Resolve the Python interpreter. Same contract as mempal_save_hook.sh: -# MEMPAL_PYTHON (explicit override) → $(command -v python3) → bare python3. -MEMPAL_PYTHON_BIN="${MEMPAL_PYTHON:-}" -if [ -z "$MEMPAL_PYTHON_BIN" ] || [ ! -x "$MEMPAL_PYTHON_BIN" ]; then - MEMPAL_PYTHON_BIN="$(command -v python3 2>/dev/null || echo python3)" -fi +# Pin to the mempalace venv (chromadb etc.). Override via MEMPAL_PYTHON. +MEMPAL_PYTHON="${MEMPAL_PYTHON:-$HOME/.mempalace/venv/bin/python}" +[ -x "$MEMPAL_PYTHON" ] || MEMPAL_PYTHON="python3" -# Read JSON input from stdin INPUT=$(cat) - -# Parse session_id and transcript_path in one call. Sanitize both, then -# read sanitized values from one-per-line stdout into shell variables — -# avoids ``eval`` on generated code (#1231 review). Same contract as -# mempal_save_hook.sh. -mapfile -t _mempal_parsed < <(echo "$INPUT" | "$MEMPAL_PYTHON_BIN" -c " -import sys, json, re -data = json.load(sys.stdin) -sid = data.get('session_id', 'unknown') -tp = data.get('transcript_path', '') -safe = lambda s: re.sub(r'[^a-zA-Z0-9_/.\-~]', '', str(s)) -print(safe(sid)) -print(safe(tp)) -" 2>/dev/null) -SESSION_ID="${_mempal_parsed[0]:-unknown}" -TRANSCRIPT_PATH="${_mempal_parsed[1]:-}" - -# Expand ~ in path -TRANSCRIPT_PATH="${TRANSCRIPT_PATH/#\~/$HOME}" - -# Validate that TRANSCRIPT_PATH looks like a transcript file. Mirrors -# mempalace.hooks_cli._validate_transcript_path so the shell hook -# rejects the same shapes the Python hook rejects (#1231 review). -is_valid_transcript_path() { - local path="$1" - [ -n "$path" ] || return 1 - case "$path" in - *.json|*.jsonl) ;; - *) return 1 ;; - esac - case "/$path/" in - */../*) return 1 ;; - esac - return 0 -} - -echo "[$(date '+%H:%M:%S')] PRE-COMPACT triggered for session $SESSION_ID" >> "$STATE_DIR/hook.log" - -# Run ingest synchronously so memories land before compaction. Two -# independent targets — both run if both are set: -# 1. TRANSCRIPT_PATH (from Claude Code) → parent dir, --mode convos -# 2. MEMPAL_DIR → --mode projects -if is_valid_transcript_path "$TRANSCRIPT_PATH" && [ -f "$TRANSCRIPT_PATH" ]; then - mempalace mine "$(dirname "$TRANSCRIPT_PATH")" --mode convos \ - >> "$STATE_DIR/hook.log" 2>&1 -elif [ -n "$TRANSCRIPT_PATH" ]; then - echo "[$(date '+%H:%M:%S')] Skipping invalid transcript path: $TRANSCRIPT_PATH" \ - >> "$STATE_DIR/hook.log" -fi -if [ -n "$MEMPAL_DIR" ] && [ -d "$MEMPAL_DIR" ]; then - mempalace mine "$MEMPAL_DIR" --mode projects \ - >> "$STATE_DIR/hook.log" 2>&1 -fi - -# Silent: return empty JSON to not block. "decision": "allow" is invalid — -# only "block" or {} are recognized. -echo '{}' +HARNESS="claude-code" +PARENT_CMD=$(ps -p $PPID -o comm= 2>/dev/null | tr -d ' ' || echo "") +case "$PARENT_CMD" in + codex|Codex) HARNESS="codex" ;; + gemini|Gemini|gemini-cli) HARNESS="gemini" ;; + qwen|Qwen|qwen-code) HARNESS="qwen" ;; + opencode|Opencode) HARNESS="opencode" ;; + *) ;; +esac + +printf '%s' "$INPUT" | "$MEMPAL_PYTHON" -m mempalace hook run --hook precompact --harness "$HARNESS" diff --git a/hooks/mempal_save_hook.sh b/hooks/mempal_save_hook.sh index 5efd157693..2a2f1df972 100755 --- a/hooks/mempal_save_hook.sh +++ b/hooks/mempal_save_hook.sh @@ -1,5 +1,5 @@ #!/bin/bash -# MEMPALACE SAVE HOOK — Auto-save every N exchanges +# MEMPALACE SAVE HOOK — Auto-detect harness (claude-code, codex, opencode, claude) # # Claude Code "Stop" hook. After every assistant response: # 1. Counts human messages in the session transcript @@ -8,9 +8,6 @@ # 4. AI does the save (topics, decisions, code, quotes → organized into palace) # 5. Next Stop fires with stop_hook_active=true → lets AI stop normally # -# The AI does the classification — it knows what wing/hall/closet to use -# because it has context about the conversation. No regex needed. -# # === INSTALL === # Add to .claude/settings.local.json: # @@ -24,200 +21,26 @@ # }] # }] # } -# -# For Codex CLI, add to .codex/hooks.json: -# -# "Stop": [{ -# "type": "command", -# "command": "/absolute/path/to/mempal_save_hook.sh", -# "timeout": 30 -# }] -# -# === HOW IT WORKS === -# -# Claude Code sends JSON on stdin with these fields: -# session_id — unique session identifier -# stop_hook_active — true if AI is already in a save cycle (prevents infinite loop) -# transcript_path — path to the JSONL transcript file -# -# When we block, Claude Code shows our "reason" to the AI as a system message. -# The AI then saves to memory, and when it tries to stop again, -# stop_hook_active=true so we let it through. No infinite loop. -# -# === MEMPALACE CLI === -# The hook ALWAYS mines the active conversation transcript automatically -# (via `mempalace mine --mode convos`). MEMPAL_DIR is an -# *additional*, optional target for project files — it does not replace -# the conversation mine. -# -# === CONFIGURATION === -SAVE_INTERVAL=15 # Save every N human messages (adjust to taste) -STATE_DIR="$HOME/.mempalace/hook_state" -mkdir -p "$STATE_DIR" +set -euo pipefail -# Optional: project directory (code / notes / docs) to also mine each -# save trigger. Mined with `--mode projects`. The conversation transcript -# is always mined regardless — this is purely additive. -# Example: MEMPAL_DIR="$HOME/projects/my_app" -MEMPAL_DIR="" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MEMPALACE_SRC="$(dirname "$SCRIPT_DIR")" +export PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$MEMPALACE_SRC" -# Resolve the Python interpreter the hook should use. -# -# Why this is nontrivial: GUI-launched Claude Code on macOS (or any harness -# that doesn't inherit the user's shell PATH) may find a `python3` on PATH -# that lacks mempalace — e.g. /usr/bin/python3 while the user installed -# mempalace into a venv or pyenv. Users in that situation can point the -# hook at the right interpreter by exporting MEMPAL_PYTHON. -# -# Resolution order (first hit wins): -# 1. $MEMPAL_PYTHON — explicit user override (absolute path) -# 2. $(command -v python3) — first python3 on the hook's PATH -# 3. bare "python3" — last-resort fallback (hope the PATH has it) -MEMPAL_PYTHON_BIN="${MEMPAL_PYTHON:-}" -if [ -z "$MEMPAL_PYTHON_BIN" ] || [ ! -x "$MEMPAL_PYTHON_BIN" ]; then - MEMPAL_PYTHON_BIN="$(command -v python3 2>/dev/null || echo python3)" -fi +# Pin to the mempalace venv (chromadb etc.). Override via MEMPAL_PYTHON. +MEMPAL_PYTHON="${MEMPAL_PYTHON:-$HOME/.mempalace/venv/bin/python}" +[ -x "$MEMPAL_PYTHON" ] || MEMPAL_PYTHON="python3" -# Read JSON input from stdin INPUT=$(cat) - -# Parse all fields in a single Python call (3x faster than separate invocations) -# without invoking ``eval`` on generated code: Python prints one sanitized -# value per line, the shell reads them via ``mapfile`` and does plain -# variable assignment — same data, smaller blast radius if the sanitizer -# is ever bypassed (#1231 review). -mapfile -t _mempal_parsed < <(echo "$INPUT" | "$MEMPAL_PYTHON_BIN" -c " -import sys, json, re -data = json.load(sys.stdin) -sid = data.get('session_id', 'unknown') -sha_raw = data.get('stop_hook_active', False) -tp = data.get('transcript_path', '') -# Shell-safe output — only allow alphanumeric, underscore, hyphen, slash, dot, tilde -safe = lambda s: re.sub(r'[^a-zA-Z0-9_/.\-~]', '', str(s)) -# Coerce stop_hook_active to strict boolean string -sha = 'True' if sha_raw is True or str(sha_raw).lower() in ('true', '1', 'yes') else 'False' -print(safe(sid)) -print(sha) -print(safe(tp)) -" 2>/dev/null) -SESSION_ID="${_mempal_parsed[0]:-unknown}" -STOP_HOOK_ACTIVE="${_mempal_parsed[1]:-False}" -TRANSCRIPT_PATH="${_mempal_parsed[2]:-}" - -# Expand ~ in path -TRANSCRIPT_PATH="${TRANSCRIPT_PATH/#\~/$HOME}" - -# Validate that TRANSCRIPT_PATH looks like a transcript file: -# - non-empty -# - .jsonl or .json suffix -# - no traversal segments (.. components) -# Mirrors mempalace.hooks_cli._validate_transcript_path so the shell hook -# rejects the same shapes the Python hook rejects (#1231 review). -is_valid_transcript_path() { - local path="$1" - [ -n "$path" ] || return 1 - case "$path" in - *.json|*.jsonl) ;; - *) return 1 ;; - esac - case "/$path/" in - */../*) return 1 ;; - esac - return 0 -} - -# If we're already in a save cycle, let the AI stop normally -# This is the infinite-loop prevention: block once → AI saves → tries to stop again → we let it through -if [ "$STOP_HOOK_ACTIVE" = "True" ] || [ "$STOP_HOOK_ACTIVE" = "true" ]; then - echo "{}" - exit 0 -fi - -# Count human messages in the JSONL transcript -# SECURITY: Pass transcript path as sys.argv to avoid shell injection via crafted paths -if [ -f "$TRANSCRIPT_PATH" ]; then - EXCHANGE_COUNT=$("$MEMPAL_PYTHON_BIN" - "$TRANSCRIPT_PATH" <<'PYEOF' -import json, sys -count = 0 -with open(sys.argv[1]) as f: - for line in f: - try: - entry = json.loads(line) - msg = entry.get('message', {}) - if isinstance(msg, dict) and msg.get('role') == 'user': - content = msg.get('content', '') - if isinstance(content, str) and '' in content: - continue - count += 1 - except: - pass -print(count) -PYEOF -2>/dev/null) -else - EXCHANGE_COUNT=0 -fi - -# Track last save point for this session -LAST_SAVE_FILE="$STATE_DIR/${SESSION_ID}_last_save" -LAST_SAVE=0 -if [ -f "$LAST_SAVE_FILE" ]; then - LAST_SAVE_RAW=$(cat "$LAST_SAVE_FILE") - # SECURITY: Validate as plain integer before arithmetic to prevent command injection - if [[ "$LAST_SAVE_RAW" =~ ^[0-9]+$ ]]; then - LAST_SAVE="$LAST_SAVE_RAW" - fi -fi - -SINCE_LAST=$((EXCHANGE_COUNT - LAST_SAVE)) - -# Log for debugging (check ~/.mempalace/hook_state/hook.log) -echo "[$(date '+%H:%M:%S')] Session $SESSION_ID: $EXCHANGE_COUNT exchanges, $SINCE_LAST since last save" >> "$STATE_DIR/hook.log" - -# Time to save? -if [ "$SINCE_LAST" -ge "$SAVE_INTERVAL" ] && [ "$EXCHANGE_COUNT" -gt 0 ]; then - # Update last save point - echo "$EXCHANGE_COUNT" > "$LAST_SAVE_FILE" - - echo "[$(date '+%H:%M:%S')] TRIGGERING SAVE at exchange $EXCHANGE_COUNT" >> "$STATE_DIR/hook.log" - - # Auto-mine. Two independent targets — both run if both are set: - # 1. TRANSCRIPT_PATH (from Claude Code) → parent dir, --mode convos - # (Claude Code session JSONL — must use the convo miner) - # 2. MEMPAL_DIR (user-configured project) → --mode projects - # (code, notes, docs) - # MEMPAL_DIR is *additive*, not an override: a user with MEMPAL_DIR - # pointed at their project still gets the active conversation mined. - if is_valid_transcript_path "$TRANSCRIPT_PATH" && [ -f "$TRANSCRIPT_PATH" ]; then - mempalace mine "$(dirname "$TRANSCRIPT_PATH")" --mode convos \ - >> "$STATE_DIR/hook.log" 2>&1 & - elif [ -n "$TRANSCRIPT_PATH" ]; then - echo "[$(date '+%H:%M:%S')] Skipping invalid transcript path: $TRANSCRIPT_PATH" \ - >> "$STATE_DIR/hook.log" - fi - if [ -n "$MEMPAL_DIR" ] && [ -d "$MEMPAL_DIR" ]; then - mempalace mine "$MEMPAL_DIR" --mode projects \ - >> "$STATE_DIR/hook.log" 2>&1 & - fi - - # MEMPAL_VERBOSE toggle: - # true = developer mode — block and show diaries/code in chat - # false = silent mode (default) — save in background, no chat clutter - # Set via: export MEMPAL_VERBOSE=true - if [ "$MEMPAL_VERBOSE" = "true" ] || [ "$MEMPAL_VERBOSE" = "1" ]; then - cat << 'HOOKJSON' -{ - "decision": "block", - "reason": "MemPalace save checkpoint. Write a brief session diary entry covering key topics, decisions, and code changes since the last save. Use verbatim quotes where possible. Continue after saving." -} -HOOKJSON - else - # Silent mode: return empty JSON to not block. "decision": "allow" is - # not a valid value — only "block" or {} are recognized. - echo '{}' - fi -else - # Not time yet — let the AI stop normally - echo "{}" -fi +HARNESS="claude-code" +PARENT_CMD=$(ps -p $PPID -o comm= 2>/dev/null | tr -d ' ' || echo "") +case "$PARENT_CMD" in + codex|Codex) HARNESS="codex" ;; + gemini|Gemini|gemini-cli) HARNESS="gemini" ;; + qwen|Qwen|qwen-code) HARNESS="qwen" ;; + opencode|Opencode) HARNESS="opencode" ;; + *) ;; +esac + +printf '%s' "$INPUT" | "$MEMPAL_PYTHON" -m mempalace hook run --hook stop --harness "$HARNESS" diff --git a/hooks/mempal_save_hook_throttled.sh b/hooks/mempal_save_hook_throttled.sh new file mode 100755 index 0000000000..2596a60f7a --- /dev/null +++ b/hooks/mempal_save_hook_throttled.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# Throttled wrapper — runs mempal_save_hook.sh every N stops (default: 10) +# Set MEMPAL_INTERVAL env var to override. + +set -euo pipefail + +INTERVAL="${MEMPAL_INTERVAL:-10}" +COUNTER_FILE="${TMPDIR:-/tmp}/mempalace_stop_counter_${UID:-0}" + +# Read current count, increment, write back +COUNT=$(cat "$COUNTER_FILE" 2>/dev/null || echo "0") +COUNT=$(( COUNT + 1 )) +echo "$COUNT" > "$COUNTER_FILE" + +if (( COUNT % INTERVAL == 0 )); then + exec "$(dirname "${BASH_SOURCE[0]}")/mempal_save_hook.sh" +else + # Must consume stdin so Claude Code hook doesn't hang + cat > /dev/null + echo "{}" +fi diff --git a/hooks/mempal_semantic_search.sh b/hooks/mempal_semantic_search.sh new file mode 100755 index 0000000000..473b9899b9 --- /dev/null +++ b/hooks/mempal_semantic_search.sh @@ -0,0 +1,76 @@ +#!/bin/bash +# MEMPALACE SEMANTIC SEARCH HELPER +# Search MemPalace for related information using semantic similarity +# +# Usage: ./mempal_semantic_search.sh "search query" [--wing wing_name] [--limit N] +# +# Examples: +# ./mempal_semantic_search.sh "project structure" +# ./mempal_semantic_search.sh "api authentication" --wing ananas-ai --limit 5 + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MEMPALACE_SRC="$(dirname "$SCRIPT_DIR")" +PALACE_PATH="/home/zapostolski/.mempalace/palace" + +# Check minimum arguments +if [[ $# -lt 1 ]]; then + echo "Error: Search query is required" + echo "Usage: $0 \"search query\" [--wing wing_name] [--limit N]" + exit 1 +fi + +QUERY="$1" +shift + +# Parse optional arguments +WING="" +LIMIT="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --wing) + WING="$2" + shift 2 + ;; + --limit) + LIMIT="$2" + shift 2 + ;; + *) + echo "Unknown option: $1" + exit 1 + ;; + esac +done + +# Create temporary Python script +TMP_PY=$(mktemp) +cat > "$TMP_PY" << EOF +import sys +sys.path.insert(0, '$MEMPALACE_SRC') +from mempalace.searcher import search_memories +import json + +query = '$QUERY' +wing = '$WING' if '$WING' else None +limit = int('$LIMIT') if '$LIMIT'.isdigit() else 5 + +try: + result = search_memories( + query, + palace_path='$PALACE_PATH', + wing=wing, + n_results=limit + ) + print(json.dumps(result, indent=2, ensure_ascii=False)) +except Exception as e: + print(json.dumps({'error': str(e)}, indent=2)) + sys.exit(1) +EOF + +python3 "$TMP_PY" +EXIT_CODE=$? +rm -f "$TMP_PY" +exit $EXIT_CODE \ No newline at end of file diff --git a/hooks/mempal_session_start_hook.sh b/hooks/mempal_session_start_hook.sh new file mode 100755 index 0000000000..485a65d293 --- /dev/null +++ b/hooks/mempal_session_start_hook.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# MEMPALACE SESSION-START HOOK +# Loads palace context (wake-up + cwd-matched wing) and a protocol nudge into the session. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MEMPALACE_SRC="$(dirname "$SCRIPT_DIR")" +export PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$MEMPALACE_SRC" + +# Pin to the mempalace venv (chromadb etc.). Override via MEMPAL_PYTHON. +MEMPAL_PYTHON="${MEMPAL_PYTHON:-$HOME/.mempalace/venv/bin/python}" +[ -x "$MEMPAL_PYTHON" ] || MEMPAL_PYTHON="python3" + +INPUT=$(cat) +HARNESS="claude-code" +PARENT_CMD=$(ps -p $PPID -o comm= 2>/dev/null | tr -d ' ' || echo "") +case "$PARENT_CMD" in + codex|Codex) HARNESS="codex" ;; + gemini|Gemini|gemini-cli) HARNESS="gemini" ;; + qwen|Qwen|qwen-code) HARNESS="qwen" ;; + opencode|Opencode) HARNESS="opencode" ;; + *) ;; +esac +printf '%s' "$INPUT" | "$MEMPAL_PYTHON" -m mempalace hook run --hook session-start --harness "$HARNESS" + +# Pre-warm ChromaDB in background so the first Stop hook doesn't cold-start +"$MEMPAL_PYTHON" -c " +import os, sys +sys.path.insert(0, os.environ.get('PYTHONPATH','').split(':')[0]) +try: + from mempalace.backends.chroma import ChromaBackend + b = ChromaBackend() + b.get_collection(os.path.expanduser('~/.mempalace/palace'), 'mempalace_drawers', create=False) +except Exception: + pass +" >/dev/null 2>&1 & diff --git a/hooks/mempal_verify_knowledge.sh b/hooks/mempal_verify_knowledge.sh new file mode 100755 index 0000000000..ca424e3b8b --- /dev/null +++ b/hooks/mempal_verify_knowledge.sh @@ -0,0 +1,81 @@ +#!/bin/bash +# MEMPALACE KNOWLEDGE VERIFICATION HELPER +# Query MemPalace knowledge graph before making statements about projects/entities +# +# Usage: ./mempal_verify_knowledge.sh "entity name" [--type relationship_type] [--limit N] +# +# Examples: +# ./mempal_verify_knowledge.sh "ananas-ai" +# ./mempal_verify_knowledge.sh "ai-marketing" --type decision --limit 5 +# ./mempal_verify_knowledge.sh "project structure" --type assessment + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MEMPALACE_SRC="$(dirname "$SCRIPT_DIR")" + +# Check minimum arguments +if [[ $# -lt 1 ]]; then + echo "Error: Entity name is required" + echo "Usage: $0 \"entity name\" [--type relationship_type] [--limit N]" + exit 1 +fi + +ENTITY="$1" +shift + +# Parse optional arguments +RELATIONSHIP_TYPE="" +LIMIT="" + +while [[ $# -gt 0 ]]; do + case "$1" in + --type) + RELATIONSHIP_TYPE="$2" + shift 2 + ;; + --limit) + LIMIT="$2" + shift 2 + ;; + *) + echo "Unknown option: $1" + exit 1 + ;; + esac +done + +# Create temporary Python script +TMP_PY=$(mktemp) +cat > "$TMP_PY" << EOF +import sys +sys.path.insert(0, '$MEMPALACE_SRC') +from mempalace.knowledge_graph import KnowledgeGraph +import json + +entity = '$ENTITY' +rel_type = '$RELATIONSHIP_TYPE' if '$RELATIONSHIP_TYPE' else None +limit = int('$LIMIT') if '$LIMIT'.isdigit() else None + +try: + kg = KnowledgeGraph() + results = kg.query_entity(entity) + + # Filter by relationship type if specified + if rel_type: + results = [r for r in results if r['predicate'] == rel_type] + + # Apply limit if specified + if limit is not None: + results = results[:limit] + + print(json.dumps(results, indent=2, ensure_ascii=False)) +except Exception as e: + print(json.dumps({'error': str(e)}, indent=2)) + sys.exit(1) +EOF + +python3 "$TMP_PY" +EXIT_CODE=$? +rm -f "$TMP_PY" +exit $EXIT_CODE \ No newline at end of file diff --git a/mempalace/backends/chroma.py b/mempalace/backends/chroma.py index 01ac62771b..bcebfec3a1 100644 --- a/mempalace/backends/chroma.py +++ b/mempalace/backends/chroma.py @@ -1,9 +1,12 @@ """ChromaDB-backed MemPalace storage backend (RFC 001 reference implementation).""" +import contextlib import datetime as _dt import logging import os +import pickle import sqlite3 +from numbers import Integral from pathlib import Path from typing import Any, Optional @@ -126,9 +129,13 @@ def _segment_appears_healthy(seg_dir: str) -> bool: return len(head) == 2 and head[0] == 0x80 and tail == b"\x2e" -def quarantine_stale_hnsw(palace_path: str, stale_seconds: float = 300.0) -> list[str]: +def quarantine_stale_hnsw( + palace_path: str, + stale_seconds: float = 300.0, + force_stale_seconds: float = 7200.0, +) -> list[str]: """Rename HNSW segment dirs that are both stale-by-mtime AND fail an - integrity sniff-test. + integrity sniff-test, or that are extremely stale regardless of metadata. Catches the segfault failure mode from #823 (semantic search stale after ``add_drawer``), observed at neo-cortex-mcp#2 (SIGSEGV on @@ -144,33 +151,36 @@ def quarantine_stale_hnsw(palace_path: str, stale_seconds: float = 300.0) -> lis 2. **Integrity gate** (``_segment_appears_healthy``). Even when the mtime gap exceeds the threshold, a segment whose - ``index_metadata.pickle`` passes a format sniff-test is healthy: - chromadb 1.5.x flushes HNSW state asynchronously and a clean - shutdown does NOT force-flush, so the on-disk HNSW is *always* - somewhat older than ``chroma.sqlite3``. Production observation - (2026-04-26 disks daemon): three of three segments quarantined - on every cold start, with 538-557s gaps, leaving the 151K-drawer - palace with vector_ranked=0 until rebuild. Renaming a healthy - segment based on mtime alone destroys a valid index — chromadb - creates an empty replacement, orphaning every drawer in sqlite - from vector recall until the operator runs ``mempalace repair - --mode rebuild`` (15+ min on a 151K palace). + ``index_metadata.pickle`` passes a format sniff-test is passed + through — chromadb 1.5.x flushes HNSW state asynchronously and a + clean shutdown does NOT force-flush, so the on-disk HNSW is + *always* somewhat older than ``chroma.sqlite3``. Production + observation (2026-04-26 disks daemon): three of three segments + quarantined on every cold start, with 538-557s gaps, leaving the + 151K-drawer palace with vector_ranked=0 until rebuild. Renaming a + healthy segment based on mtime alone destroys a valid index. + + **Exception — force quarantine.** If the mtime gap exceeds + ``force_stale_seconds`` (default 2 h), the segment is quarantined + regardless of metadata health. ChromaDB's async flush-lag is + measured in seconds, never hours. A gap of 2+ hours indicates the + segment was written by a previous process that no longer holds an + open handle; the metadata sniff-test cannot detect all binary + incompatibilities between ChromaDB versions that cause SIGSEGV on + load (observed with chromadb 0.6.x loading segments written under + earlier versions). Only segments that pass stage 1 (suspiciously stale) AND fail stage - 2 (metadata file truncated, zero-filled, or absent-with-data) are - renamed to ``.drift-``. The original directory is - renamed, not deleted, so recovery remains possible if the heuristic - misfires. - - The default threshold (5 min) is advisory under daemon-strict; the - integrity gate is what actually distinguishes corruption from flush - lag. The threshold still matters for the cross-machine replication - case (#823), where it bounds how stale a Syncthing-replicated - segment can be before we look harder at it. + 2 (metadata file truncated, zero-filled, or absent-with-data), or + that exceed ``force_stale_seconds``, are renamed to + ``.drift-``. The original directory is renamed, not + deleted, so recovery remains possible if the heuristic misfires. Args: palace_path: path to the palace directory containing ``chroma.sqlite3`` stale_seconds: minimum mtime gap to *consider* a segment for quarantine + force_stale_seconds: mtime gap above which quarantine is forced even + when the metadata sniff-test passes (default 2 h) Returns: List of paths that were quarantined (empty if nothing actually @@ -208,18 +218,31 @@ def quarantine_stale_hnsw(palace_path: str, stale_seconds: float = 300.0) -> lis # Stage 2: integrity gate. mtime drift is necessary but not # sufficient — chromadb's async flush makes drift the steady- - # state condition. A healthy segment metadata file proves - # chromadb can open the segment without segfault; don't - # quarantine a healthy index. - if _segment_appears_healthy(seg_dir): + # state condition. A healthy segment metadata file normally + # proves chromadb can open the segment without segfault. + # Exception: very large gaps (>force_stale_seconds) are forced + # through regardless — flush-lag is never measured in hours, and + # chromadb version mismatches can cause SIGSEGV even on metadata- + # healthy segments when the gap is this large. + gap = sqlite_mtime - hnsw_mtime + if _segment_appears_healthy(seg_dir) and gap < force_stale_seconds: logger.info( "HNSW mtime gap %.0fs on %s exceeds threshold but segment " "metadata file is intact — flush-lag, not corruption. " "Leaving in place.", - sqlite_mtime - hnsw_mtime, + gap, seg_dir, ) continue + if _segment_appears_healthy(seg_dir): + logger.warning( + "HNSW mtime gap %.0fs on %s exceeds force-quarantine threshold " + "(%.0fs) — quarantining despite intact metadata to prevent " + "potential SIGSEGV on chromadb version mismatch.", + gap, + seg_dir, + force_stale_seconds, + ) stamp = _dt.datetime.now().strftime("%Y%m%d-%H%M%S") target = f"{seg_dir}.drift-{stamp}" @@ -489,22 +512,17 @@ def hnsw_capacity_status(palace_path: str, collection_name: str = "mempalace_dra divergence_floor = max(_HNSW_DIVERGENCE_FALLBACK_FLOOR, 2 * sync_threshold) if hnsw_count is None: - # No pickle yet — segment hasn't persisted metadata. Could be - # fresh-but-unflushed (normal) or interrupted-mid-flush (bad). - # We can't distinguish without the pickle, so only flag - # divergence when sqlite holds clearly more than two flush - # windows worth — same threshold as the with-pickle path. - if sqlite_count > divergence_floor: - out["status"] = "diverged" - out["diverged"] = True - out["divergence"] = sqlite_count - out["message"] = ( - f"sqlite holds {sqlite_count:,} embeddings but the HNSW segment " - "has never flushed metadata — vector search will return nothing " - "until the segment is rebuilt. Run `mempalace repair`." - ) - else: - out["message"] = "HNSW segment metadata not yet flushed; skipping" + # No pickle yet, so this probe cannot measure HNSW capacity. + # Chroma 1.5.x can have binary HNSW files without a flushed + # metadata pickle; absence of the pickle alone is not proof that + # vector search is unusable or dangerous. Keep the status unknown + # so MCP does not globally disable vectors on an inconclusive + # signal. Corrupt/invalid metadata, when present, is handled by + # quarantine_invalid_hnsw_metadata before Chroma opens. + out["message"] = ( + "HNSW capacity unavailable: metadata has not been flushed; " + "leaving vector search enabled" + ) return out divergence = sqlite_count - hnsw_count @@ -591,6 +609,97 @@ def _pin_hnsw_threads(collection) -> None: _BLOB_FIX_MARKER = ".blob_seq_ids_migrated" +def _valid_dimensionality(value: object) -> bool: + return isinstance(value, Integral) and not isinstance(value, bool) and int(value) > 0 + + +def _persisted_metadata_fields(obj: object) -> tuple[object, object]: + if isinstance(obj, dict): + return obj.get("dimensionality"), obj.get("id_to_label") + return getattr(obj, "dimensionality", None), getattr(obj, "id_to_label", None) + + +def quarantine_invalid_hnsw_metadata(palace_path: str) -> list[str]: + """Quarantine segment dirs whose ``index_metadata.pickle`` is unreadable or invalid. + + Chroma's persisted HNSW metadata is untrusted disk state. If a segment has + labels but no valid positive dimensionality, current Chroma versions can + accept the pickle and crash later in the Rust loader. We rename the entire + segment out of the way before ``PersistentClient`` opens so Chroma can + rebuild cleanly instead of touching known-bad metadata. + """ + try: + entries = os.listdir(palace_path) + except OSError: + return [] + + moved: list[str] = [] + for name in entries: + if "-" not in name or name.startswith(".") or ".drift-" in name or ".corrupt-" in name: + continue + seg_dir = os.path.join(palace_path, name) + if not os.path.isdir(seg_dir): + continue + + meta_path = os.path.join(seg_dir, "index_metadata.pickle") + if not os.path.isfile(meta_path): + continue + + reason = None + try: + persisted = _SafePersistentDataUnpickler.load(meta_path) + except (EOFError, OSError): + logger.debug( + "Skipping invalid-HNSW quarantine for transient metadata read in %s", + meta_path, + exc_info=True, + ) + continue + except pickle.UnpicklingError as exc: + if "truncated" in str(exc).lower() or "ran out of input" in str(exc).lower(): + logger.debug( + "Skipping invalid-HNSW quarantine for transient metadata read in %s", + meta_path, + exc_info=True, + ) + continue + reason = f"invalid index_metadata.pickle: {exc}" + except Exception as exc: + reason = f"invalid index_metadata.pickle: {exc}" + else: + if not isinstance(persisted, dict) and not ( + hasattr(persisted, "dimensionality") or hasattr(persisted, "id_to_label") + ): + reason = f"unrecognized index_metadata.pickle payload: {type(persisted).__name__}" + else: + dimensionality, id_to_label = _persisted_metadata_fields(persisted) + if id_to_label is not None and not isinstance(id_to_label, dict): + reason = f"invalid id_to_label type {type(id_to_label).__name__}" + else: + has_labels = bool(id_to_label) + if has_labels and not _valid_dimensionality(dimensionality): + reason = ( + "labels present but dimensionality is missing or invalid " + f"({dimensionality!r})" + ) + elif dimensionality is not None and not _valid_dimensionality(dimensionality): + reason = f"invalid dimensionality {dimensionality!r}" + + if reason is None: + continue + + stamp = _dt.datetime.now().strftime("%Y%m%d-%H%M%S") + target = f"{seg_dir}.corrupt-{stamp}" + try: + os.rename(seg_dir, target) + moved.append(target) + logger.warning("Quarantined invalid HNSW metadata in %s: %s", seg_dir, reason) + except OSError: + logger.exception("Failed to quarantine invalid HNSW metadata in %s", seg_dir) + + return moved + + def _fix_blob_seq_ids(palace_path: str) -> None: """Fix ChromaDB 0.6.x -> 1.5.x migration bug: BLOB seq_ids -> INTEGER. @@ -677,10 +786,43 @@ def _as_list(v: Any) -> list: class ChromaCollection(BaseCollection): - """Thin adapter translating ChromaDB dict returns into typed results.""" + """Thin adapter translating ChromaDB dict returns into typed results. + + When ``palace_path`` is set, all write methods (``add``, ``upsert``, + ``update``, ``delete``) acquire ``mine_palace_lock(palace_path)`` for the + duration of the underlying chromadb call. This serializes MCP and other + direct-backend writers against ``mempalace mine`` and against each other, + closing the race between concurrent writers that triggers ChromaDB's + multi-threaded HNSW corruption (#974/#965). + + The lock is the same primitive used by ``miner.mine()`` so re-entrant + acquisition from inside the mine pipeline (mine -> _mine_body -> + collection.upsert) is short-circuited by the per-thread guard inside + ``mine_palace_lock`` — no self-deadlock. + + ``palace_path=None`` disables the wrapping, preserving the legacy + no-lock behaviour for callers that construct a ``ChromaCollection`` + directly without going through ``ChromaBackend``. + """ - def __init__(self, collection): + def __init__(self, collection, palace_path: Optional[str] = None): self._collection = collection + self._palace_path = palace_path + + @contextlib.contextmanager + def _write_lock(self): + """Acquire ``mine_palace_lock`` for the configured palace, if any. + + No-op (yields immediately) when ``self._palace_path`` is None. + """ + if self._palace_path is None: + yield + return + # Late import — palace.py imports ChromaBackend from this module. + from ..palace import mine_palace_lock + + with mine_palace_lock(self._palace_path): + yield # ------------------------------------------------------------------ # Writes @@ -692,7 +834,8 @@ def add(self, *, documents, ids, metadatas=None, embeddings=None): kwargs["metadatas"] = metadatas if embeddings is not None: kwargs["embeddings"] = embeddings - self._collection.add(**kwargs) + with self._write_lock(): + self._collection.add(**kwargs) def upsert(self, *, documents, ids, metadatas=None, embeddings=None): kwargs: dict[str, Any] = {"documents": documents, "ids": ids} @@ -700,7 +843,8 @@ def upsert(self, *, documents, ids, metadatas=None, embeddings=None): kwargs["metadatas"] = metadatas if embeddings is not None: kwargs["embeddings"] = embeddings - self._collection.upsert(**kwargs) + with self._write_lock(): + self._collection.upsert(**kwargs) def update( self, @@ -719,7 +863,8 @@ def update( kwargs["metadatas"] = metadatas if embeddings is not None: kwargs["embeddings"] = embeddings - self._collection.update(**kwargs) + with self._write_lock(): + self._collection.update(**kwargs) # ------------------------------------------------------------------ # Reads @@ -863,7 +1008,8 @@ def delete(self, *, ids=None, where=None): kwargs["ids"] = ids if where is not None: kwargs["where"] = where - self._collection.delete(**kwargs) + with self._write_lock(): + self._collection.delete(**kwargs) def count(self): return self._collection.count() @@ -994,6 +1140,12 @@ def _client(self, palace_path: str): if cached is None or inode_changed or mtime_changed or mtime_appeared: _fix_blob_seq_ids(palace_path) + if inode_changed: + ChromaBackend._quarantined_paths.discard(palace_path) + if palace_path not in ChromaBackend._quarantined_paths: + quarantine_invalid_hnsw_metadata(palace_path) + quarantine_stale_hnsw(palace_path) + ChromaBackend._quarantined_paths.add(palace_path) cached = chromadb.PersistentClient(path=palace_path) self._clients[palace_path] = cached # Re-stat after the client constructor runs: chromadb creates @@ -1006,26 +1158,27 @@ def _client(self, palace_path: str): # Public static helpers (legacy; prefer :meth:`get_collection`) # ------------------------------------------------------------------ - # Per-process record of palaces that have already had quarantine_stale_hnsw - # invoked at least once. The proactive drift check is a *cold-start* - # protection — it catches HNSW segments that arrived stale relative to - # ``chroma.sqlite3`` (e.g. cross-machine replication, partial restore, - # crashed-mid-write). Once a long-running process has opened the palace - # cleanly, re-firing on every reconnect is a *runtime thrash*: the - # daemon's own writes bump sqlite mtime but HNSW flushes batch on - # chromadb's internal cadence, so the mtime gap naturally exceeds the - # threshold under steady write load even though nothing is corrupt. + # Per-process record of palaces that have already had the cold-start + # quarantine invoked at least once. The proactive HNSW checks are a + # *cold-start* protection — they catch segments that arrive stale relative + # to ``chroma.sqlite3`` or invalid on disk (e.g. cross-machine replication, + # partial restore, crashed-mid-write). Once a long-running process has + # opened the palace cleanly, re-firing the stale check on every reconnect + # is a *runtime thrash*: the daemon's own writes bump sqlite mtime but HNSW + # flushes batch on chromadb's internal cadence, so the mtime gap naturally + # exceeds the threshold under steady write load even though nothing is + # corrupt. # Real runtime drift is still handled — palace-daemon's ``_auto_repair`` # calls :func:`quarantine_stale_hnsw` directly on observed HNSW errors, # which bypasses this gate. # # Thread-safety: this set is mutated without a lock. Two concurrent # ``make_client()`` calls for the same palace can both pass the - # membership check and both invoke ``quarantine_stale_hnsw``. That's - # safe because the function is idempotent (mtime check + timestamped - # rename of distinct directories), so the worst-case race produces - # one redundant rename attempt that no-ops. Idempotency is the - # safety property; locking would add cost without correctness gain. + # membership check and both invoke the cold-start quarantine. That's + # safe because the functions are idempotent (mtime checks + timestamped + # rename of distinct directories), so the worst-case race produces one + # redundant rename attempt that no-ops. Idempotency is the safety + # property; locking would add cost without correctness gain. _quarantined_paths: set[str] = set() @staticmethod @@ -1036,12 +1189,13 @@ def make_client(palace_path: str): own client cache. New code should obtain a collection through :meth:`get_collection` which manages caching internally. - Quarantines stale HNSW segments **once per palace per process**. See + Quarantines HNSW segments **once per palace per process**. See :attr:`_quarantined_paths` for the rationale (cold-start protection vs. runtime thrash on steady-write daemons). """ _fix_blob_seq_ids(palace_path) if palace_path not in ChromaBackend._quarantined_paths: + quarantine_invalid_hnsw_metadata(palace_path) quarantine_stale_hnsw(palace_path) ChromaBackend._quarantined_paths.add(palace_path) return chromadb.PersistentClient(path=palace_path) @@ -1109,7 +1263,7 @@ def get_collection( else: collection = client.get_collection(collection_name, **ef_kwargs) _pin_hnsw_threads(collection) - return ChromaCollection(collection) + return ChromaCollection(collection, palace_path=palace_path) def close_palace(self, palace) -> None: """Drop cached handles for ``palace``. Accepts ``PalaceRef`` or legacy path str.""" @@ -1160,7 +1314,7 @@ def create_collection( }, **ef_kwargs, ) - return ChromaCollection(collection) + return ChromaCollection(collection, palace_path=palace_path) def _normalize_get_collection_args(args, kwargs): diff --git a/mempalace/cli.py b/mempalace/cli.py index ca9798b444..c5bf030d8e 100644 --- a/mempalace/cli.py +++ b/mempalace/cli.py @@ -310,8 +310,7 @@ def cmd_init(args): ) except LLMError as e: print( - f" LLM init failed ({e}). " - f"Running heuristics-only — pass --no-llm to silence this." + f" LLM init failed ({e}). Running heuristics-only — pass --no-llm to silence this." ) # Pass 0: detect whether the corpus is AI-dialogue. Writes @@ -505,6 +504,7 @@ def cmd_mine(args): limit=args.limit, dry_run=args.dry_run, extract_mode=args.extract, + include_subagents=args.include_subagents, ) else: from .miner import mine @@ -648,10 +648,19 @@ def cmd_repair(args): import shutil from .backends.chroma import ChromaBackend from .migrate import confirm_destructive_action, contains_palace_database - from .repair import TruncationDetected, check_extraction_safety + from .repair import ( + RebuildCollectionError, + TruncationDetected, + _close_chroma_handles, + _extract_drawers, + _rebuild_collection_via_temp, + check_extraction_safety, + ) + config = MempalaceConfig() + collection_name = config.collection_name palace_path = os.path.abspath( - os.path.expanduser(args.palace) if args.palace else MempalaceConfig().palace_path + os.path.expanduser(args.palace) if args.palace else config.palace_path ) if getattr(args, "mode", "legacy") == "max-seq-id": @@ -685,7 +694,7 @@ def cmd_repair(args): # Try to read existing drawers try: - col = backend.get_collection(palace_path, "mempalace_drawers") + col = backend.get_collection(palace_path, collection_name) total = col.count() print(f" Drawers found: {total}") except Exception as e: @@ -705,18 +714,7 @@ def cmd_repair(args): # Extract all drawers in batches print("\n Extracting drawers...") batch_size = 5000 - all_ids = [] - all_docs = [] - all_metas = [] - offset = 0 - while offset < total: - batch = col.get(limit=batch_size, offset=offset, include=["documents", "metadatas"]) - if not batch["ids"]: - break - all_ids.extend(batch["ids"]) - all_docs.extend(batch["documents"]) - all_metas.extend(batch["metadatas"]) - offset += len(batch["ids"]) + all_ids, all_docs, all_metas = _extract_drawers(col, total, batch_size) print(f" Extracted {len(all_ids)} drawers") # ── #1208 guard ────────────────────────────────────────────────── @@ -731,12 +729,12 @@ def cmd_repair(args): palace_path, len(all_ids), confirm_truncation_ok=getattr(args, "confirm_truncation_ok", False), + collection_name=collection_name, ) except TruncationDetected as e: print(e.message) return - # Backup and rebuild palace_path = os.path.normpath(palace_path) backup_path = palace_path + ".backup" if os.path.exists(backup_path): @@ -750,18 +748,34 @@ def cmd_repair(args): print(f" Backing up to {backup_path}...") shutil.copytree(palace_path, backup_path) - print(" Rebuilding collection...") - backend.delete_collection(palace_path, "mempalace_drawers") - new_col = backend.create_collection(palace_path, "mempalace_drawers") - - filed = 0 - for i in range(0, len(all_ids), batch_size): - batch_ids = all_ids[i : i + batch_size] - batch_docs = all_docs[i : i + batch_size] - batch_metas = all_metas[i : i + batch_size] - new_col.add(documents=batch_docs, ids=batch_ids, metadatas=batch_metas) - filed += len(batch_ids) - print(f" Re-filed {filed}/{len(all_ids)} drawers...") + try: + filed = _rebuild_collection_via_temp( + backend, + palace_path, + all_ids, + all_docs, + all_metas, + batch_size, + collection_name=collection_name, + progress=print, + ) + except RebuildCollectionError as e: + print(f" Repair failed: {e}") + if getattr(e, "live_replaced", False): + print(" Live collection was already replaced; restoring from backup...") + try: + _close_chroma_handles(palace_path, backend=backend) + if os.path.exists(palace_path): + shutil.rmtree(palace_path) + shutil.copytree(backup_path, palace_path) + print(f" Restore complete from backup: {backup_path}") + except Exception as restore_error: + print(f" Automatic restore failed: {restore_error}") + print(" Manual recovery required:") + print(f" 1. Remove or rename the broken directory: {palace_path}") + print(f" 2. Restore the backup directory to: {palace_path}") + print(f" Backup location: {backup_path}") + sys.exit(1) print(f"\n Repair complete. {filed} drawers rebuilt.") print(f" Backup saved at {backup_path}") @@ -902,6 +916,15 @@ def cmd_compress(args): # Store compressed versions (unless dry-run) if not args.dry_run: try: + # Drop and recreate the compressed collection on each run. + # Repeated upserts in chromadb 1.5.8 cause the HNSW link_lists.bin + # sparse file to grow without GC; rebuilding the index from scratch + # each compress run keeps disk size proportional to entry count. + # See #1092 for the broader concurrent-writer report. + try: + backend.delete_collection(palace_path, "mempalace_compressed") + except Exception: + pass comp_col = backend.get_or_create_collection(palace_path, "mempalace_compressed") for doc_id, compressed, meta, stats in compressed_entries: comp_meta = dict(meta) @@ -1079,6 +1102,17 @@ def main(): default="exchange", help="Extraction strategy for convos mode: 'exchange' (default) or 'general' (5 memory types)", ) + p_mine.add_argument( + "--include-subagents", + action="store_true", + default=False, + help=( + "Also mine Claude Code subagent transcripts (subagents/ dirs). " + "Excluded by default: these are short ephemeral exchanges " + "(Explore/Plan/Grep agents) already summarized in the parent " + "session, and on typical workspaces they dominate file counts." + ), + ) # sweep p_sweep = sub.add_parser( @@ -1153,7 +1187,7 @@ def main(): p_hook_run.add_argument( "--harness", required=True, - choices=["claude-code", "codex"], + choices=["claude-code", "codex", "opencode", "gemini", "qwen", "deepseek"], help="Harness type (determines stdin JSON format)", ) diff --git a/mempalace/config.py b/mempalace/config.py index cacd1f9184..3afab3f288 100644 --- a/mempalace/config.py +++ b/mempalace/config.py @@ -4,10 +4,14 @@ Priority: env vars > config file (~/.mempalace/config.json) > defaults """ +from __future__ import annotations + import json import os import re +from functools import lru_cache from pathlib import Path +from typing import Optional # ── Input validation ────────────────────────────────────────────────────────── @@ -92,9 +96,43 @@ def sanitize_content(value: str, max_length: int = 100_000) -> str: return value +# ── ISO-8601 date validation ───────────────────────────────────────────────── +# Accepts YYYY, YYYY-MM, or YYYY-MM-DD. Used at the MCP boundary so that +# invalid date strings are rejected early instead of silently producing +# empty knowledge-graph query results. + +_ISO_DATE_RE = re.compile(r"^\d{4}(?:-(?:0[1-9]|1[0-2])(?:-(?:0[1-9]|[12]\d|3[01]))?)?$") + + +def validate_iso_date(value: Optional[str], param_name: str = "date") -> Optional[str]: + """Validate an optional ISO-8601 date string (YYYY, YYYY-MM, or YYYY-MM-DD). + + Returns the value unchanged if valid (or ``None``). Raises ``ValueError`` + with a user-facing message if the format is unrecognised. + """ + if value is None: + return None + if not isinstance(value, str) or not value.strip(): + return None + value = value.strip() + if not _ISO_DATE_RE.match(value): + raise ValueError( + f"{param_name}={value!r} is not a valid ISO-8601 date " + f"(expected YYYY, YYYY-MM, or YYYY-MM-DD)" + ) + return value + + DEFAULT_PALACE_PATH = os.path.expanduser("~/.mempalace/palace") DEFAULT_COLLECTION_NAME = "mempalace_drawers" + +@lru_cache(maxsize=1) +def get_configured_collection_name() -> str: + """Return the configured drawer collection name without repeated config-file reads.""" + return MempalaceConfig().collection_name + + DEFAULT_TOPIC_WINGS = [ "emotions", "consciousness", diff --git a/mempalace/convo_miner.py b/mempalace/convo_miner.py index 2cf57e4488..2cb2714a3c 100644 --- a/mempalace/convo_miner.py +++ b/mempalace/convo_miner.py @@ -279,12 +279,26 @@ def detect_convo_room(content: str) -> str: # ============================================================================= -def scan_convos(convo_dir: str) -> list: - """Find all potential conversation files.""" +def scan_convos(convo_dir: str, include_subagents: bool = False) -> list: + """Find all potential conversation files. + + By default, directories named ``subagents`` are skipped: Claude Code + records Explore/Plan/Grep subagent transcripts there, and on typical + workspaces they outnumber main session files by one to two orders of + magnitude. Pass ``include_subagents=True`` to mine them anyway. + + The match is case-insensitive on the directory name only (``subagents`` + or ``Subagents``), so directories like ``mysubagents`` or + ``subagentsbackup`` are not affected. + """ convo_path = Path(convo_dir).expanduser().resolve() files = [] for root, dirs, filenames in os.walk(convo_path): - dirs[:] = [d for d in dirs if d not in SKIP_DIRS] + dirs[:] = [ + d + for d in dirs + if d not in SKIP_DIRS and (include_subagents or d.lower() != "subagents") + ] for filename in filenames: if filename.endswith(".meta.json"): continue @@ -384,12 +398,16 @@ def mine_convos( limit: int = 0, dry_run: bool = False, extract_mode: str = "exchange", + include_subagents: bool = False, ): """Mine a directory of conversation files into the palace. extract_mode: "exchange" — default exchange-pair chunking (Q+A = one unit) "general" — general extractor: decisions, preferences, milestones, problems, emotions + include_subagents: + False (default) — skip Claude Code ``subagents/`` directories + True — also mine subagent transcripts """ convo_path = Path(convo_dir).expanduser().resolve() @@ -398,7 +416,7 @@ def mine_convos( wing = normalize_wing_name(convo_path.name) - files = scan_convos(convo_dir) + files = scan_convos(convo_dir, include_subagents=include_subagents) if limit > 0: files = files[:limit] diff --git a/mempalace/entity_registry.py b/mempalace/entity_registry.py index 78d8a8b20c..c8ac517249 100644 --- a/mempalace/entity_registry.py +++ b/mempalace/entity_registry.py @@ -16,6 +16,7 @@ """ import json +import os import re import urllib.request import urllib.parse @@ -320,11 +321,35 @@ def save(self): self._path.parent.chmod(0o700) except (OSError, NotImplementedError): pass - self._path.write_text(json.dumps(self._data, indent=2), encoding="utf-8") + # Atomic write: serialize to a sibling temp file in the same dir + # (so os.replace stays on one filesystem), fsync, then rename over + # the target. A crash mid-write leaves the previous registry intact + # instead of a half-written file or an empty file from the truncate. + payload = json.dumps(self._data, indent=2) + tmp_path = self._path.with_name(self._path.name + ".tmp") + with open(tmp_path, "w", encoding="utf-8") as f: + f.write(payload) + f.flush() + os.fsync(f.fileno()) try: - self._path.chmod(0o600) + tmp_path.chmod(0o600) except (OSError, NotImplementedError): pass + os.replace(tmp_path, self._path) + # On ext4 (and similar) the rename's durability across power loss + # requires an additional fsync on the parent directory. Without it, + # the kernel can ack the rename and a crash reverts to the state + # where the temp file is present and the target is at the old version. + try: + dir_fd = os.open(str(self._path.parent), os.O_RDONLY) + try: + os.fsync(dir_fd) + finally: + os.close(dir_fd) + except OSError: + # Windows and some special filesystems reject directory fds — they + # have different durability semantics on rename anyway. + pass @staticmethod def _empty() -> dict: diff --git a/mempalace/hooks_cli.py b/mempalace/hooks_cli.py index d4f8317232..7fe5435152 100644 --- a/mempalace/hooks_cli.py +++ b/mempalace/hooks_cli.py @@ -14,8 +14,19 @@ from datetime import datetime from pathlib import Path -SAVE_INTERVAL = 15 +SAVE_INTERVAL = 10 STATE_DIR = Path.home() / ".mempalace" / "hook_state" +MIRROR_STATE_FILE = STATE_DIR / "mirror_state.json" + +# Local memory dirs for each agent harness — scanned on Stop hook and mirrored +# into mempalace as drawers. Override via MEMPAL_MIRROR_ROOTS env (colon-separated). +DEFAULT_MIRROR_ROOTS = ( + Path.home() / ".claude" / "projects", + Path.home() / ".codex" / "memories", + Path.home() / ".gemini" / "memory", + Path.home() / ".qwen" / "memory", +) +MIRROR_SKIP_FILENAMES = {"MEMORY.md", "CLAUDE.md", "GEMINI.md", "QWEN.md", "AGENTS.md"} def _mempalace_python() -> str: @@ -94,46 +105,146 @@ def _validate_transcript_path(transcript_path: str) -> Path: return path -def _count_human_messages(transcript_path: str) -> int: - """Count human messages in a JSONL transcript, skipping command-messages.""" +def _normalize_text(text: str) -> str: + """Normalize transcript text for lightweight heuristics.""" + return re.sub(r"\s+", " ", text).strip() + + +def _extract_text(content) -> str: + """Flatten transcript content blocks into plain text.""" + if isinstance(content, str): + return content + if isinstance(content, list): + blocks = [] + for block in content: + if isinstance(block, dict): + text = block.get("text", "") + if isinstance(text, str): + blocks.append(text) + return " ".join(blocks) + return "" + + +def _iter_real_messages(transcript_path: str): + """Yield normalized user/assistant messages, excluding command chatter.""" path = _validate_transcript_path(transcript_path) - if path is None: - if transcript_path: - _log(f"WARNING: transcript_path rejected by validator: {transcript_path!r}") - return 0 - if not path.is_file(): - return 0 - count = 0 + if path is None or not path.is_file(): + return try: with open(path, encoding="utf-8", errors="replace") as f: for line in f: try: entry = json.loads(line) msg = entry.get("message", {}) - if isinstance(msg, dict) and msg.get("role") == "user": - content = msg.get("content", "") - if isinstance(content, str): - if "" in content: - continue - elif isinstance(content, list): - text = " ".join( - b.get("text", "") for b in content if isinstance(b, dict) - ) - if "" in text: - continue - count += 1 - # Also handle Codex CLI transcript format - # {"type": "event_msg", "payload": {"type": "user_message", "message": "..."}} - elif entry.get("type") == "event_msg": - payload = entry.get("payload", {}) - if isinstance(payload, dict) and payload.get("type") == "user_message": - msg_text = payload.get("message", "") - if isinstance(msg_text, str) and "" not in msg_text: - count += 1 except (json.JSONDecodeError, AttributeError): - pass + continue + if not isinstance(msg, dict): + continue + role = msg.get("role") + if role not in {"user", "assistant"}: + continue + text = _normalize_text(_extract_text(msg.get("content", ""))) + if not text or "" in text: + continue + yield role, text except OSError: - return 0 + return + + +def _iter_real_messages_any(transcript_path: str): + """Yield user/assistant messages from ANY transcript format. + + Tries multiple schemas in order: + 1. Claude Code JSONL: {"message": {"role": ..., "content": ...}} + 2. Qwen JSONL: {"message": {"role": ..., "parts": [{"text": ...}]}} + 3. JSON files (Gemini, Claude sessions): via normalize.py + 4. normalize.py fallback for any format + """ + path = _validate_transcript_path(transcript_path) + if path is None or not path.is_file(): + return + + ext = path.suffix.lower() + + # Try Qwen JSONL schema: {"message": {"role": "...", "parts": [{"text": "..."}]}} + if ext == ".jsonl": + try: + found = False + with open(path, encoding="utf-8", errors="replace") as f: + for line in f: + try: + entry = json.loads(line) + except (json.JSONDecodeError, AttributeError): + continue + msg = entry.get("message", {}) + if not isinstance(msg, dict): + continue + role = msg.get("role") + if role not in {"user", "assistant"}: + continue + parts = msg.get("parts", []) + text = "" + if isinstance(parts, list): + text = " ".join(p.get("text", "") for p in parts if isinstance(p, dict)) + if text and "" not in text: + found = True + yield role, _normalize_text(text) + if found: + return + except OSError: + pass + + # Try standard JSONL schema (Claude Code, Codex) + for role, text in _iter_real_messages(str(path)): + yield role, text + + # For JSON files (Gemini, Claude session files), use normalize.py + if ext == ".json": + try: + from mempalace.normalize import normalize + transcript = normalize(str(path)) + if transcript and "> " in transcript: + lines = transcript.split("\n") + i = 0 + while i < len(lines): + if lines[i].startswith("> "): + yield "user", lines[i][2:].strip() + i += 1 + if i < len(lines) and lines[i].strip() and not lines[i].startswith("> "): + yield "assistant", lines[i].strip() + i += 1 + else: + i += 1 + return + except Exception: + pass + + # Final fallback: normalize.py for any format + try: + from mempalace.normalize import normalize + transcript = normalize(str(path)) + if transcript and "> " in transcript: + lines = transcript.split("\n") + i = 0 + while i < len(lines): + if lines[i].startswith("> "): + yield "user", lines[i][2:].strip() + i += 1 + if i < len(lines) and lines[i].strip() and not lines[i].startswith("> "): + yield "assistant", lines[i].strip() + i += 1 + else: + i += 1 + except Exception: + pass + + +def _count_human_messages(transcript_path: str) -> int: + """Count human messages in any transcript format, skipping command chatter.""" + count = 0 + for role, _text in _iter_real_messages_any(transcript_path) or []: + if role == "user": + count += 1 return count @@ -337,41 +448,11 @@ def _desktop_toast(body: str, title: str = "MemPalace"): def _extract_recent_messages(transcript_path: str, count: int = _RECENT_MSG_COUNT) -> list[str]: - """Extract the last N user messages from a JSONL transcript.""" - path = Path(transcript_path).expanduser() - if not path.is_file(): - return [] + """Extract the last N user messages from any transcript format.""" messages = [] - try: - with open(path, encoding="utf-8", errors="replace") as f: - for line in f: - try: - entry = json.loads(line) - # Claude Code format - msg = entry.get("message") or entry.get("event_message") or {} - if isinstance(msg, dict) and msg.get("role") == "user": - content = msg.get("content", "") - if isinstance(content, list): - content = " ".join( - b.get("text", "") for b in content if isinstance(b, dict) - ) - if not isinstance(content, str) or not content.strip(): - continue - if "" in content or "" in content: - continue - messages.append(content.strip()[:200]) - # Codex CLI format - elif entry.get("type") == "event_msg": - payload = entry.get("payload", {}) - if isinstance(payload, dict) and payload.get("type") == "user_message": - text = payload.get("message", "") - if isinstance(text, str) and text.strip(): - if "" not in text: - messages.append(text.strip()[:200]) - except (json.JSONDecodeError, AttributeError): - pass - except OSError: - return [] + for role, text in _iter_real_messages_any(transcript_path) or []: + if role == "user" and text.strip(): + messages.append(text.strip()[:200]) return messages[-count:] @@ -465,8 +546,12 @@ def _save_diary_direct( def _ingest_transcript(transcript_path: str): """Mine a Claude Code session transcript into the palace as a conversation.""" - path = Path(transcript_path).expanduser() - if not path.is_file() or path.stat().st_size < 100: + path = _validate_transcript_path(transcript_path) + if path is None or not path.is_file() or path.stat().st_size < 100: + return + + if _mine_already_running(): + _log(f"Skipping transcript ingest: mine already running") return from .config import MempalaceConfig @@ -480,7 +565,7 @@ def _ingest_transcript(transcript_path: str): log_path = STATE_DIR / "hook.log" STATE_DIR.mkdir(parents=True, exist_ok=True) with open(log_path, "a") as log_f: - subprocess.Popen( + proc = subprocess.Popen( [ _mempalace_python(), "-m", @@ -495,12 +580,138 @@ def _ingest_transcript(transcript_path: str): stdout=log_f, stderr=log_f, ) - _log(f"Transcript ingest started: {path.name}") + _MINE_PID_FILE.write_text(str(proc.pid)) + _log(f"Transcript ingest started: {path.name} (pid={proc.pid})") + except OSError: + pass + + +def _maybe_sync_obsidian(): + """Mirror recent memory changes into the Obsidian vault when available.""" + sync_script = Path.home() / "obsidian-vault" / "sync.py" + if not sync_script.is_file(): + return + try: + log_path = STATE_DIR / "hook.log" + with open(log_path, "a") as log_f: + subprocess.Popen( + [sys.executable, str(sync_script), "--quick"], + stdout=log_f, + stderr=log_f, + ) except OSError: pass -SUPPORTED_HARNESSES = {"claude-code", "codex"} +SUPPORTED_HARNESSES = {"claude-code", "claude", "codex", "opencode", "gemini", "qwen", "deepseek"} +SIGNAL_KEYWORDS = ( + "decision", + "plan", + "fix", + "build", + "implement", + "found", + "finding", + "audit", + "error", + "blocker", + "risk", + "deploy", + "architecture", + "infra", + "database", + "metric", + "source", + "query", + "report", + "powerbi", + "portal", + "aws", + "rds", + "ecs", + "mcp", + "api", + "route", + "component", + "dataset", + "dax", + "migration", + "verify", + "passed", + "failed", +) +ARTIFACT_HINT_RE = re.compile( + r"`[^`]+`|/[\w./-]+|\b[\w.-]+\.(?:ts|tsx|js|jsx|py|sh|md|sql|json|ya?ml|tf)\b|https?://", + re.IGNORECASE, +) +TRIVIAL_MESSAGE_RE = re.compile( + r"^(?:" + r"continue|resume|go on|keep going|proceed|" + r"ok(?:ay)?|k|yes|yep|no|nah|" + r"thanks|thank you|good|sounds good|do it|go ahead|" + r"continue please|switch and continue|" + r"codex resume --last" + r")(?:[.!? ]+)?$", + re.IGNORECASE, +) + + +def _collect_unsaved_messages(transcript_path: str, last_save: int) -> list[str]: + """Collect the message slice after the last checkpointed user message.""" + messages: list[str] = [] + user_count = 0 + for role, text in _iter_real_messages_any(transcript_path) or []: + if role == "user": + user_count += 1 + if user_count <= last_save: + continue + elif user_count <= last_save: + continue + messages.append(text) + return messages + + +def _looks_trivial(text: str) -> bool: + """Treat short acknowledgements as low signal.""" + normalized = _normalize_text(text).lower() + if not normalized: + return True + if TRIVIAL_MESSAGE_RE.fullmatch(normalized): + return True + return len(normalized) < 8 + + +def _has_meaningful_updates(messages: list[str]) -> bool: + """Avoid blocking on chatter-only windows.""" + substantive: list[str] = [] + keyword_hits = 0 + artifact_hits = 0 + + for text in messages: + if _looks_trivial(text): + continue + lower = text.lower() + has_keyword = any(keyword in lower for keyword in SIGNAL_KEYWORDS) + artifact_count = len(list(ARTIFACT_HINT_RE.finditer(text))) + if has_keyword: + keyword_hits += 1 + artifact_hits += artifact_count + if len(text) >= 30 or has_keyword or artifact_count: + substantive.append(text) + + if not substantive: + return False + + char_count = sum(len(text) for text in substantive) + if keyword_hits >= 2: + return True + if artifact_hits >= 2: + return True + if len(substantive) >= 4 and char_count >= 300: + return True + if len(substantive) >= 2 and char_count >= 500: + return True + return False def _parse_harness_input(data: dict, harness: str) -> dict: @@ -512,6 +723,8 @@ def _parse_harness_input(data: dict, harness: str) -> dict: "session_id": _sanitize_session_id(str(data.get("session_id", "unknown"))), "stop_hook_active": data.get("stop_hook_active", False), "transcript_path": str(data.get("transcript_path", "")), + "cwd": str(data.get("cwd", "")), + "source": str(data.get("source", "")), } @@ -548,6 +761,23 @@ def _wing_from_transcript_path(transcript_path: str) -> str: return "wing_sessions" +def _auto_kg_write(session_id: str, themes: list[str]) -> None: + """Add session-to-topic KG relationships for extracted themes.""" + try: + from .mcp_server import tool_kg_add + date_str = datetime.now().strftime("%Y-%m-%d") + for theme in themes[:3]: + tool_kg_add( + subject=f"session-{date_str}", + predicate="covered_topic", + object=theme, + valid_from=date_str, + ) + _log(f"KG auto-write: session-{date_str} → {themes[:3]}") + except Exception as exc: + _log(f"KG auto-write failed: {exc}") + + def hook_stop(data: dict, harness: str): """Stop hook: block every N messages for auto-save.""" parsed = _parse_harness_input(data, harness) @@ -555,6 +785,17 @@ def hook_stop(data: dict, harness: str): stop_hook_active = parsed["stop_hook_active"] transcript_path = parsed["transcript_path"] + # Always mirror local memory dirs first — runs even when we don't block, + # so .md files Claude writes get auto-replicated into mempalace drawers + # regardless of whether the LLM cooperated with the save protocol. + if os.environ.get("MEMPAL_MIRROR_DISABLED", "").lower() not in ("1", "true", "yes"): + try: + counts = _mirror_local_memory() + if counts.get("added") or counts.get("errors"): + _log(f"mirror: {counts}") + except Exception as exc: + _log(f"mirror: unexpected failure ({exc}); continuing") + # If already in a block-mode save cycle, let through (infinite-loop prevention). # Silent mode saves directly without returning {"decision":"block"}, so there's # no loop to prevent — and Claude Code's plugin dispatch sets this flag on every @@ -597,10 +838,20 @@ def hook_stop(data: dict, harness: str): _log(f"Session {session_id}: {exchange_count} exchanges, {since_last} since last save") if since_last >= SAVE_INTERVAL and exchange_count > 0: + unsaved_messages = _collect_unsaved_messages(transcript_path, last_save) + if not _has_meaningful_updates(unsaved_messages): + _log(f"SKIPPING SAVE at exchange {exchange_count}: low-signal window") + _output({}) + return + _log(f"TRIGGERING SAVE at exchange {exchange_count}") # Read hook settings from config from .config import MempalaceConfig + + # Optional: auto-ingest if MEMPAL_DIR is set + _maybe_auto_ingest() + _maybe_sync_obsidian() try: config = MempalaceConfig() @@ -613,6 +864,14 @@ def hook_stop(data: dict, harness: str): project_wing = _wing_from_transcript_path(transcript_path) if silent: + # Skip diary save if a mine subprocess is running — concurrent + # ChromaDB writes cause "another mine already running" failures. + # Don't advance last_save so next Stop hook retries. + if _mine_already_running(): + _log("Skipping diary save: mine running. Will retry on next stop.") + _output({}) + return + # Save directly via Python API — systemMessage renders in terminal result = {"count": 0} if transcript_path: @@ -620,7 +879,6 @@ def hook_stop(data: dict, harness: str): transcript_path, session_id, wing=project_wing, toast=toast ) _ingest_transcript(transcript_path) - _maybe_auto_ingest() # Only advance save marker after successful save count = result.get("count", 0) if count > 0: @@ -633,9 +891,24 @@ def hook_stop(data: dict, harness: str): tag = " \u2014 " + ", ".join(themes) else: tag = "" + + # Auto-write session topics to KG + if themes: + _auto_kg_write(session_id, themes) + + # Include recent KG facts in system message to keep agent context fresh. + try: + from .layers import MemoryStack + from .config import MempalaceConfig + stack = MemoryStack(palace_path=MempalaceConfig().palace_path) + kg_summary = stack.lkg.generate(limit=5) + kg_note = f"\n\n{kg_summary}" if "No current" not in kg_summary else "" + except Exception: + kg_note = "" + _output( { - "systemMessage": f"\u2726 {count} memories woven into the palace{tag}", + "systemMessage": f"\u2726 {count} memories woven into the palace{tag}{kg_note}", } ) else: @@ -650,25 +923,240 @@ def hook_stop(data: dict, harness: str): pass if transcript_path: _ingest_transcript(transcript_path) - _maybe_auto_ingest() reason = STOP_BLOCK_REASON + f" Write diary entry to wing={project_wing}." _output({"decision": "block", "reason": reason}) else: _output({}) +# ============================================================================= +# LOCAL MEMORY MIRROR — replicate per-agent .md memory dirs into mempalace +# ============================================================================= + + +def _mirror_roots() -> list[Path]: + """Return list of memory roots to scan, honoring MEMPAL_MIRROR_ROOTS env.""" + override = os.environ.get("MEMPAL_MIRROR_ROOTS", "").strip() + if override: + return [Path(p).expanduser() for p in override.split(":") if p.strip()] + return list(DEFAULT_MIRROR_ROOTS) + + +def _load_mirror_state() -> dict: + """Read the {abs_path: mtime} map of files already mirrored.""" + if not MIRROR_STATE_FILE.is_file(): + return {} + try: + with open(MIRROR_STATE_FILE, encoding="utf-8") as f: + data = json.load(f) + return data if isinstance(data, dict) else {} + except (OSError, json.JSONDecodeError): + return {} + + +def _save_mirror_state(state: dict) -> None: + """Persist the {abs_path: mtime} map atomically.""" + try: + STATE_DIR.mkdir(parents=True, exist_ok=True) + tmp = MIRROR_STATE_FILE.with_suffix(".json.tmp") + with open(tmp, "w", encoding="utf-8") as f: + json.dump(state, f, indent=2) + os.replace(tmp, MIRROR_STATE_FILE) + except OSError as exc: + _log(f"mirror_state write failed: {exc}") + + +def _derive_wing_room(file_path: Path, root: Path) -> tuple[str, str]: + """Derive (wing, room) from a memory file's path relative to its root. + + Conventions: + - Claude: ~/.claude/projects//memory/.md + slug like '-home-zapostolski-projects-ai-marketing' → wing 'ai-marketing' + - Gemini: ~/.gemini/memory//.md + - Codex/Qwen: ~/.codex/memories//.md or flat ~/.codex/memories/.md + + Falls back to: + wing = first directory under root, or 'misc' + room = filename stem + """ + try: + rel = file_path.relative_to(root) + except ValueError: + return "misc", file_path.stem + parts = rel.parts + room = file_path.stem + + if not parts: + return "misc", room + + first = parts[0] + # Claude project slug: -home-...-projects- + if "-projects-" in first: + wing = first.rsplit("-projects-", 1)[1] + else: + # Strip leading dashes from agent slugs + wing = first.lstrip("-") or "misc" + return wing, room + + +def _mirror_local_memory() -> dict: + """Mirror new/changed .md files from local memory dirs into mempalace drawers. + + Idempotent: tracks (path, mtime) state in MIRROR_STATE_FILE. Catches all + errors so a mirror failure never blocks the harness. Returns counts for logging. + """ + counts = {"scanned": 0, "added": 0, "skipped_dup": 0, "skipped_unchanged": 0, "errors": 0} + try: + from .mcp_server import tool_add_drawer + except Exception as exc: + _log(f"mirror: cannot import tool_add_drawer ({exc}); skipping") + return counts + + state = _load_mirror_state() + new_state = dict(state) + + for root in _mirror_roots(): + if not root.is_dir(): + continue + for md_path in root.rglob("*.md"): + if md_path.name in MIRROR_SKIP_FILENAMES: + continue + counts["scanned"] += 1 + try: + mtime = md_path.stat().st_mtime + except OSError: + counts["errors"] += 1 + continue + key = str(md_path) + if state.get(key) == mtime: + counts["skipped_unchanged"] += 1 + continue + try: + content = md_path.read_text(encoding="utf-8", errors="replace") + except OSError: + counts["errors"] += 1 + continue + if not content.strip(): + new_state[key] = mtime + continue + wing, room = _derive_wing_room(md_path, root) + try: + result = tool_add_drawer( + wing=wing, + room=room, + content=content, + source_file=str(md_path), + added_by="auto-mirror", + ) + except Exception as exc: + _log(f"mirror: add_drawer raised for {md_path}: {exc}") + counts["errors"] += 1 + continue + if result.get("success"): + counts["added"] += 1 + new_state[key] = mtime + elif result.get("reason") == "duplicate": + counts["skipped_dup"] += 1 + # Mark as seen so we don't keep re-checking unchanged dupes + new_state[key] = mtime + else: + counts["errors"] += 1 + _log(f"mirror: add_drawer failed for {md_path}: {result}") + + if new_state != state: + _save_mirror_state(new_state) + return counts + + +def _wing_from_cwd(cwd: str) -> str: + """Best-effort match of cwd to a known palace wing. + + Strategy: take cwd's basename and check it against the wings reported by + `tool_status()`. Wings live in metadata, not on disk, so a directory check + is unreliable. + """ + if not cwd: + return "" + candidate = os.path.basename(cwd.rstrip("/")) + if not candidate: + return "" + try: + from .mcp_server import tool_status + wings = tool_status().get("wings", {}) or {} + except Exception: + return "" + return candidate if candidate in wings else "" + + +PROTOCOL_NUDGE = ( + "MemPalace protocol: BEFORE responding about projects/people/decisions, " + "call mempalace_kg_query or mempalace_search to verify. WHEN making " + "decisions/plans/architecture changes, call mempalace_add_drawer " + "(room='decisions' or 'plans') and mempalace_kg_add for the relationships. " + "AFTER each session, the Stop hook will prompt you to checkpoint — route " + "those into the palace, not local files." +) + + +def _build_session_start_context(cwd: str, palace_path: str) -> str: + """Build SessionStart additionalContext: wake-up text + protocol nudge. + + Wing match is based on the cwd's basename matching a palace wing directory. + """ + sections: list[str] = [] + wing = _wing_from_cwd(cwd) + try: + from .layers import MemoryStack + stack = MemoryStack(palace_path=palace_path) + wake_text = stack.wake_up(wing=wing) if wing else stack.wake_up() + if wake_text: + header = ( + f"# MemPalace wake-up (wing={wing})" + if wing else "# MemPalace wake-up" + ) + sections.append(f"{header}\n\n{wake_text}") + except Exception as exc: + sections.append(f"# MemPalace wake-up unavailable: {exc}") + sections.append(PROTOCOL_NUDGE) + return "\n\n".join(sections) + + +def _wrap_session_start_output(harness: str, context: str) -> dict: + """Format additionalContext per harness expectations. + + Claude Code uses hookSpecificOutput.additionalContext. Other harnesses + that mimic Claude's schema accept the same shape; harnesses that don't + will see a no-op (the unknown key is ignored). + """ + if not context: + return {} + return { + "hookSpecificOutput": { + "hookEventName": "SessionStart", + "additionalContext": context, + } + } + + def hook_session_start(data: dict, harness: str): - """Session start hook: initialize session tracking state.""" + """Session start hook: inject palace context (status + wing match + protocol nudge).""" parsed = _parse_harness_input(data, harness) session_id = parsed["session_id"] + cwd = parsed["cwd"] - _log(f"SESSION START for session {session_id}") - - # Initialize session state directory + _log(f"SESSION START for session {session_id} cwd={cwd!r} harness={harness}") STATE_DIR.mkdir(parents=True, exist_ok=True) - # Pass through — no blocking on session start - _output({}) + try: + from .config import MempalaceConfig + palace_path = MempalaceConfig().palace_path + except Exception as exc: + _log(f"SessionStart: cannot resolve palace_path ({exc}); skipping context inject") + _output({}) + return + + context = _build_session_start_context(cwd, palace_path) + _output(_wrap_session_start_output(harness, context)) def hook_precompact(data: dict, harness: str): @@ -683,12 +1171,32 @@ def hook_precompact(data: dict, harness: str): if transcript_path: _ingest_transcript(transcript_path) - # Mine MEMPAL_DIR synchronously so project data lands before - # compaction proceeds. Transcript convos were already kicked off - # above via _ingest_transcript. + # Mine MEMPAL_DIR synchronously _mine_sync() - _output({}) + # KG Integration: Check behavior mode for blocking + # Supported: block (always), block_once (per session), proceed (never) + mode = os.environ.get("MEMPAL_PRECOMPACT_MODE", "").lower() + if not mode: + # block_once per session for all harnesses: saves before compaction + # without risking infinite loops (flag is per session_id). + mode = "block_once" + + if mode == "proceed": + _output({}) + return + + # Check session-once state if needed + if mode == "block_once": + STATE_DIR.mkdir(parents=True, exist_ok=True) + flag = STATE_DIR / f"{session_id}_precompact_blocked" + if flag.exists(): + _output({}) + return + flag.touch() + + # Block and force manual KG/Drawer save + _output({"decision": "block", "reason": PRECOMPACT_BLOCK_REASON}) def run_hook(hook_name: str, harness: str): diff --git a/mempalace/knowledge_graph.py b/mempalace/knowledge_graph.py index 9096ab28f7..59f37d7411 100644 --- a/mempalace/knowledge_graph.py +++ b/mempalace/knowledge_graph.py @@ -171,6 +171,15 @@ def add_triple( add_triple("Max", "does", "swimming", valid_from="2025-01-01") add_triple("Alice", "worried_about", "Max injury", valid_from="2026-01", valid_to="2026-02") """ + # Reject inverted intervals: a triple with valid_to < valid_from + # would never satisfy `valid_from <= as_of AND valid_to >= as_of`, + # so it would be invisible to every query — silently corrupt. + if valid_from is not None and valid_to is not None and valid_to < valid_from: + raise ValueError( + f"valid_to={valid_to!r} is before valid_from={valid_from!r}; " + "an inverted interval would be invisible to every KG query" + ) + sub_id = self._entity_id(subject) obj_id = self._entity_id(obj) pred = predicate.lower().replace(" ", "_") @@ -390,6 +399,34 @@ def stats(self): "relationship_types": predicates, } + def get_summary(self, limit: int = 20) -> str: + """Return a compact text summary of the most recent N current facts.""" + with self._lock: + conn = self._conn() + rows = conn.execute( + """ + SELECT s.name as sub_name, t.predicate, o.name as obj_name, t.valid_from + FROM triples t + JOIN entities s ON t.subject = s.id + JOIN entities o ON t.object = o.id + WHERE t.valid_to IS NULL + ORDER BY t.extracted_at DESC + LIMIT ? + """, + (limit,), + ).fetchall() + + if not rows: + return "No current knowledge graph facts." + + lines = ["## KG — CURRENT RELATIONSHIPS"] + for r in rows: + line = f" - {r['sub_name']} \u2192 {r['predicate']} \u2192 {r['obj_name']}" + if r["valid_from"]: + line += f" (since {r['valid_from']})" + lines.append(line) + return "\n".join(lines) + # ── Seed from known facts ───────────────────────────────────────────── def seed_from_entity_facts(self, entity_facts: dict): diff --git a/mempalace/layers.py b/mempalace/layers.py index a0f9b6df0d..9d2119ee6e 100644 --- a/mempalace/layers.py +++ b/mempalace/layers.py @@ -24,6 +24,7 @@ from .config import MempalaceConfig from .palace import get_collection as _get_collection from .searcher import _first_or_empty, build_where_filter +from .knowledge_graph import KnowledgeGraph # --------------------------------------------------------------------------- @@ -177,6 +178,26 @@ def generate(self) -> str: return "\n".join(lines) +# --------------------------------------------------------------------------- +# Layer KG — Knowledge Graph Summary +# --------------------------------------------------------------------------- + + +class LayerKG: + """ + Compact summary of current KG facts. + """ + + def __init__(self, palace_path: str = None): + # KG path is typically a sibling to the palace dir + cfg = MempalaceConfig() + db_path = os.path.join(palace_path or cfg.palace_path, "knowledge_graph.sqlite3") + self.kg = KnowledgeGraph(db_path=db_path) + + def generate(self, limit: int = 15) -> str: + return self.kg.get_summary(limit=limit) + + # --------------------------------------------------------------------------- # Layer 2 — On-Demand (wing/room filtered retrieval) # --------------------------------------------------------------------------- @@ -371,12 +392,13 @@ def __init__(self, palace_path: str = None, identity_path: str = None): self.l0 = Layer0(self.identity_path) self.l1 = Layer1(self.palace_path) + self.lkg = LayerKG(self.palace_path) self.l2 = Layer2(self.palace_path) self.l3 = Layer3(self.palace_path) def wake_up(self, wing: str = None) -> str: """ - Generate wake-up text: L0 (identity) + L1 (essential story). + Generate wake-up text: L0 (identity) + L1 (essential story) + KG summary. Typically ~600-900 tokens. Inject into system prompt or first message. Args: @@ -392,6 +414,10 @@ def wake_up(self, wing: str = None) -> str: if wing: self.l1.wing = wing parts.append(self.l1.generate()) + parts.append("") + + # KG Summary + parts.append(self.lkg.generate()) return "\n".join(parts) @@ -415,6 +441,9 @@ def status(self) -> dict: "L1_essential": { "description": "Auto-generated from top palace drawers", }, + "LKG_knowledge_graph": { + "description": "Compact summary of current KG relationships", + }, "L2_on_demand": { "description": "Wing/room filtered retrieval", }, @@ -445,7 +474,7 @@ def usage(): print("layers.py — 4-Layer Memory Stack") print() print("Usage:") - print(" python layers.py wake-up Show L0 + L1") + print(" python layers.py wake-up Show L0 + L1 + KG") print(" python layers.py wake-up --wing=NAME Wake-up for a specific project") print(" python layers.py recall --wing=NAME On-demand L2 retrieval") print(" python layers.py search Deep L3 search") diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 7d6211a0b4..bcf5d245cc 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -47,7 +47,7 @@ import logging # noqa: E402 import hashlib # noqa: E402 import time # noqa: E402 -from datetime import datetime # noqa: E402 +from datetime import date, datetime # noqa: E402 from pathlib import Path # noqa: E402 from typing import Optional # noqa: E402 @@ -56,6 +56,7 @@ sanitize_kg_value, sanitize_name, sanitize_content, + validate_iso_date, ) from .version import __version__ # noqa: E402 from chromadb.errors import NotFoundError as _ChromaNotFoundError # noqa: E402 @@ -67,6 +68,7 @@ _pin_hnsw_threads, hnsw_capacity_status, ) +from .embedding import get_embedding_function # noqa: E402 from .query_sanitizer import sanitize_query # noqa: E402 from .searcher import search_memories # noqa: E402 from .palace_graph import ( # noqa: E402 @@ -92,6 +94,7 @@ def _parse_args(): metavar="PATH", help="Path to the palace directory (overrides config file and env var)", ) + # Ignore unrelated argv flags so imports work cleanly under test runners. args, unknown = parser.parse_known_args() if unknown: logger.debug("Ignoring unknown args: %s", unknown) @@ -104,12 +107,10 @@ def _parse_args(): os.environ["MEMPALACE_PALACE_PATH"] = os.path.abspath(_args.palace) _config = MempalaceConfig() -# Only override KG path when --palace is explicitly provided; otherwise use -# KnowledgeGraph's default (~/.mempalace/knowledge_graph.sqlite3). -if _args.palace: - _kg = KnowledgeGraph(db_path=os.path.join(_config.palace_path, "knowledge_graph.sqlite3")) -else: - _kg = KnowledgeGraph() +# Always co-locate the KG with the palace so layers.py (wake-up) and MCP +# tools read/write the same database. The old default (~/.mempalace/knowledge_graph.sqlite3) +# caused a split where kg_add wrote to root KG but session wake-up read palace KG. +_kg = KnowledgeGraph(db_path=os.path.join(_config.palace_path, "knowledge_graph.sqlite3")) _client_cache = None @@ -142,7 +143,7 @@ def _refresh_vector_disabled_flag() -> None: """ global _vector_disabled, _vector_disabled_reason, _vector_capacity_status try: - info = hnsw_capacity_status(_config.palace_path, "mempalace_drawers") + info = hnsw_capacity_status(_config.palace_path, _config.collection_name) except Exception: logger.debug("HNSW capacity probe raised", exc_info=True) return @@ -279,6 +280,22 @@ def _get_collection(create=False): global _collection_cache, _metadata_cache, _metadata_cache_time try: client = _get_client() + # ChromaDB 1.x does not persist the embedding function with the + # collection, so a reader/writer that omits ``embedding_function=`` + # silently gets the chromadb-built-in default. On bleeding-edge + # interpreters (#1299: python 3.14 + chromadb 1.5.x on Apple Silicon) + # the default's lazy ONNX provider selection can SIGSEGV the host + # process on first ``col.add()``. The miner / Stop hook ingest path + # avoids this because it routes through ``ChromaBackend.get_collection`` + # which resolves the EF via ``mempalace.embedding.get_embedding_function``. + # The MCP server bypassed that abstraction; mirror its behaviour so + # ``tool_diary_write`` / ``tool_add_drawer`` get the same EF as mining. + try: + ef = get_embedding_function() + except Exception: + logger.exception("Failed to build embedding function; using chromadb default") + ef = None + ef_kwargs = {"embedding_function": ef} if ef is not None else {} if create: # hnsw:num_threads=1 disables ChromaDB's multi-threaded ParallelFor # HNSW insert path, which has a race in repairConnectionsForUpdate / @@ -293,7 +310,7 @@ def _get_collection(create=False): # below skips the metadata-comparison codepath for existing # collections, mirroring the backend-layer fix from #1262. try: - raw = client.get_collection(_config.collection_name) + raw = client.get_collection(_config.collection_name, **ef_kwargs) except _ChromaNotFoundError: raw = client.create_collection( _config.collection_name, @@ -302,15 +319,16 @@ def _get_collection(create=False): "hnsw:num_threads": 1, **_HNSW_BLOAT_GUARD, }, + **ef_kwargs, ) _pin_hnsw_threads(raw) - _collection_cache = ChromaCollection(raw) + _collection_cache = ChromaCollection(raw, palace_path=_config.palace_path) _metadata_cache = None _metadata_cache_time = 0 elif _collection_cache is None: - raw = client.get_collection(_config.collection_name) + raw = client.get_collection(_config.collection_name, **ef_kwargs) _pin_hnsw_threads(raw) - _collection_cache = ChromaCollection(raw) + _collection_cache = ChromaCollection(raw, palace_path=_config.palace_path) _metadata_cache = None _metadata_cache_time = 0 return _collection_cache @@ -392,6 +410,7 @@ def _tool_status_via_sqlite() -> dict: db_path = os.path.join(_config.palace_path, "chroma.sqlite3") if not os.path.isfile(db_path): return _no_palace() + collection_name = _config.collection_name wings: dict = {} rooms: dict = {} @@ -405,8 +424,9 @@ def _tool_status_via_sqlite() -> dict: FROM embeddings e JOIN segments s ON e.segment_id = s.id JOIN collections c ON s.collection = c.id - WHERE c.name = 'mempalace_drawers' - """ + WHERE c.name = ? + """, + (collection_name,), ).fetchone() total = int(row[0]) if row and row[0] is not None else 0 for key, target in (("wing", wings), ("room", rooms)): @@ -417,12 +437,12 @@ def _tool_status_via_sqlite() -> dict: JOIN embeddings e ON em.id = e.id JOIN segments s ON e.segment_id = s.id JOIN collections c ON s.collection = c.id - WHERE c.name = 'mempalace_drawers' + WHERE c.name = ? AND em.key = ? AND em.string_value IS NOT NULL GROUP BY em.string_value """, - (key,), + (collection_name, key), ): target[value] = count finally: @@ -624,6 +644,7 @@ def tool_search( n_results=limit, max_distance=dist, vector_disabled=_vector_disabled, + collection_name=_config.collection_name, ) if _vector_disabled: result["vector_disabled"] = True @@ -825,8 +846,8 @@ def tool_add_drawer( # Idempotency: if the deterministic ID already exists, return success as a no-op. try: - existing = col.get(ids=[drawer_id]) - if existing and existing["ids"]: + existing = col.get(ids=[drawer_id], include=[]) + if existing.ids: return {"success": True, "reason": "already_exists", "drawer_id": drawer_id} except Exception: pass @@ -846,6 +867,12 @@ def tool_add_drawer( } ], ) + inserted = col.get(ids=[drawer_id], include=[]) + if not inserted.ids: + raise RuntimeError( + "Drawer write was acknowledged but the new ID is not readable. " + "The palace index may be stale; run reconnect or repair." + ) _metadata_cache = None logger.info(f"Filed drawer: {drawer_id} → {wing}/{room}") return {"success": True, "drawer_id": drawer_id, "wing": wing, "room": room} @@ -1033,6 +1060,7 @@ def tool_kg_query(entity: str, as_of: str = None, direction: str = "both"): """Query the knowledge graph for an entity's relationships.""" try: entity = sanitize_kg_value(entity, "entity") + as_of = validate_iso_date(as_of, "as_of") except ValueError as e: return {"error": str(e)} if direction not in ("outgoing", "incoming", "both"): @@ -1042,13 +1070,31 @@ def tool_kg_query(entity: str, as_of: str = None, direction: str = "both"): def tool_kg_add( - subject: str, predicate: str, object: str, valid_from: str = None, source_closet: str = None + subject: str, + predicate: str, + object: str, + valid_from: str = None, + valid_to: str = None, + source_closet: str = None, + source_file: str = None, + source_drawer_id: str = None, ): - """Add a relationship to the knowledge graph.""" + """Add a relationship to the knowledge graph. + + All temporal and provenance fields are optional. ``valid_to`` lets callers + backfill historical facts with a known end date in a single call (instead + of a separate ``kg_invalidate``). ``source_file`` and ``source_drawer_id`` + are RFC 002 §5.5 provenance fields populated by adapters / bulk importers. + + TODO(#1283): once the ISO-8601 validation PR lands, wire ``validate_iso_date`` + over ``valid_from`` / ``valid_to`` here so malformed dates fail fast at the + MCP boundary instead of silently producing empty query results. + """ try: subject = sanitize_kg_value(subject, "subject") predicate = sanitize_name(predicate, "predicate") object = sanitize_kg_value(object, "object") + valid_from = validate_iso_date(valid_from, "valid_from") except ValueError as e: return {"success": False, "error": str(e)} @@ -1059,32 +1105,57 @@ def tool_kg_add( "predicate": predicate, "object": object, "valid_from": valid_from, + "valid_to": valid_to, "source_closet": source_closet, + "source_file": source_file, + "source_drawer_id": source_drawer_id, }, ) triple_id = _kg.add_triple( - subject, predicate, object, valid_from=valid_from, source_closet=source_closet + subject, + predicate, + object, + valid_from=valid_from, + valid_to=valid_to, + source_closet=source_closet, + source_file=source_file, + source_drawer_id=source_drawer_id, ) return {"success": True, "triple_id": triple_id, "fact": f"{subject} → {predicate} → {object}"} def tool_kg_invalidate(subject: str, predicate: str, object: str, ended: str = None): - """Mark a fact as no longer true (set end date).""" + """Mark a fact as no longer true (set end date). + + Returns the actual ``ended`` date that was stored — when the caller omits + ``ended``, the underlying graph stamps ``date.today()``, and the response + reflects that resolved value (instead of the literal string ``"today"``) + so callers can verify what was persisted. + + TODO(#1283): apply ``validate_iso_date`` to ``ended`` once that PR lands. + """ try: subject = sanitize_kg_value(subject, "subject") predicate = sanitize_name(predicate, "predicate") object = sanitize_kg_value(object, "object") + ended = validate_iso_date(ended, "ended") except ValueError as e: return {"success": False, "error": str(e)} + resolved_ended = ended or date.today().isoformat() _wal_log( "kg_invalidate", - {"subject": subject, "predicate": predicate, "object": object, "ended": ended}, + { + "subject": subject, + "predicate": predicate, + "object": object, + "ended": resolved_ended, + }, ) - _kg.invalidate(subject, predicate, object, ended=ended) + _kg.invalidate(subject, predicate, object, ended=resolved_ended) return { "success": True, "fact": f"{subject} → {predicate} → {object}", - "ended": ended or "today", + "ended": resolved_ended, } @@ -1114,9 +1185,13 @@ def tool_diary_write(agent_name: str, entry: str, topic: str = "general", wing: This is the agent's personal journal — observations, thoughts, what it worked on, what it noticed, what it thinks matters. + + Note: ``agent_name`` is normalized to lowercase before storage so + that diary reads are case-insensitive (see #1243). "Claude", + "claude", and "CLAUDE" all resolve to the same agent. """ try: - agent_name = sanitize_name(agent_name, "agent_name") + agent_name = sanitize_name(agent_name, "agent_name").lower() entry = sanitize_content(entry) topic = sanitize_name(topic, "topic") except ValueError as e: @@ -1125,7 +1200,7 @@ def tool_diary_write(agent_name: str, entry: str, topic: str = "general", wing: if wing: wing = sanitize_name(wing) else: - wing = f"wing_{agent_name.lower().replace(' ', '_')}" + wing = f"wing_{agent_name.replace(' ', '_')}" room = "diary" col = _get_collection(create=True) if not col: @@ -1190,9 +1265,14 @@ def tool_diary_read(agent_name: str, last_n: int = 10, wing: str = ""): written to. Diary writes from hooks land in project-derived wings (``wing_``), so requiring a specific wing on read would silo those entries from agent-initiated reads. + + Note: ``agent_name`` is normalized to lowercase before filtering so + that reads are case-insensitive (see #1243). Entries written under + pre-fix mixed-case agent names will not match the lowercase filter; + use ``mempalace repair`` to migrate legacy data if needed. """ try: - agent_name = sanitize_name(agent_name, "agent_name") + agent_name = sanitize_name(agent_name, "agent_name").lower() if wing: wing = sanitize_name(wing) except ValueError as e: @@ -1336,6 +1416,30 @@ def tool_reconnect(): _palace_db_mtime, \ _vector_disabled, \ _vector_disabled_reason + from . import palace as palace_module + + close_errors = [] + try: + palace_module._DEFAULT_BACKEND.close_palace(_config.palace_path) + except Exception as exc: + logger.debug("Failed to close shared palace backend during reconnect", exc_info=True) + close_errors.append(f"backend close_palace failed: {exc}") + try: + from chromadb.api.client import SharedSystemClient + + clear_system_cache = getattr(SharedSystemClient, "clear_system_cache", None) + if callable(clear_system_cache): + clear_system_cache() + else: + logger.debug( + "SharedSystemClient.clear_system_cache is unavailable; skipping shared Chroma cache clear during reconnect" + ) + except Exception as exc: + logger.debug( + "Failed to clear Chroma shared system cache during reconnect", + exc_info=True, + ) + close_errors.append(f"shared Chroma cache clear failed: {exc}") _client_cache = None _collection_cache = None _palace_db_inode = 0 @@ -1348,12 +1452,24 @@ def tool_reconnect(): try: col = _get_collection() if col is None: - return { + result = { "success": False, "message": "No palace found after reconnect", "drawers": 0, "vector_disabled": _vector_disabled, } + if close_errors: + result["error"] = "; ".join(close_errors) + return result + if close_errors: + return { + "success": False, + "message": "Reconnect reopened the palace but failed to fully reset cached handles", + "drawers": col.count(), + "vector_disabled": _vector_disabled, + "vector_disabled_reason": _vector_disabled_reason, + "error": "; ".join(close_errors), + } return { "success": True, "message": "Reconnected to palace", @@ -1421,7 +1537,7 @@ def tool_reconnect(): "handler": tool_kg_query, }, "mempalace_kg_add": { - "description": "Add a fact to the knowledge graph. Subject → predicate → object with optional time window. E.g. ('Max', 'started_school', 'Year 7', valid_from='2026-09-01').", + "description": "Add a fact to the knowledge graph. Subject → predicate → object with optional time window. E.g. ('Max', 'started_school', 'Year 7', valid_from='2026-09-01'). Pass valid_to to backfill an already-ended historical fact in a single call.", "input_schema": { "type": "object", "properties": { @@ -1435,10 +1551,22 @@ def tool_reconnect(): "type": "string", "description": "When this became true (YYYY-MM-DD, optional)", }, + "valid_to": { + "type": "string", + "description": "When this stopped being true (YYYY-MM-DD, optional). Use for backfilling already-ended historical facts.", + }, "source_closet": { "type": "string", "description": "Closet ID where this fact appears (optional)", }, + "source_file": { + "type": "string", + "description": "Source file path the fact was extracted from (optional)", + }, + "source_drawer_id": { + "type": "string", + "description": "Drawer ID the fact was extracted from (optional, RFC 002 §5.5 provenance)", + }, }, "required": ["subject", "predicate", "object"], }, @@ -1819,7 +1947,7 @@ def handle_request(request): "id": req_id, "result": { "protocolVersion": negotiated, - "capabilities": {"tools": {}}, + "capabilities": {"tools": {}, "prompts": {}}, "serverInfo": {"name": "mempalace", "version": __version__}, }, } @@ -1898,9 +2026,60 @@ def handle_request(request): "error": {"code": -32000, "message": "Internal tool error"}, } + elif method == "prompts/list": + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "prompts": [ + { + "name": "mempalace_checkpoint", + "description": ( + "Auto-save rule: every 10 user messages, save new decisions, " + "findings, milestones, and next steps to mempalace using " + "mempalace_add_drawer. Skip chatter and duplicates." + ), + } + ] + }, + } + elif method == "prompts/get": + name = params.get("name", "") + if name == "mempalace_checkpoint": + return { + "jsonrpc": "2.0", + "id": req_id, + "result": { + "description": "Auto-save checkpoint rule for mempalace", + "messages": [ + { + "role": "user", + "content": { + "type": "text", + "text": ( + "MEMPALACE AUTO-SAVE RULE (always active):\n" + "Every 10 user messages, before replying, call mempalace_add_drawer " + "to save any new decisions, findings, milestones, blockers, " + "architecture changes, code contracts, or actionable next steps " + "from this conversation. Skip duplicates and chatter. " + "Use the appropriate wing/room based on content type. " + "After saving, continue with your normal reply." + ), + }, + } + ], + }, + } + return { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32602, "message": f"Unknown prompt: {name}"}, + } + # Notifications (missing id) must never get a response if req_id is None: return None + return { "jsonrpc": "2.0", "id": req_id, @@ -1939,12 +2118,11 @@ def main(): request = json.loads(line) response = handle_request(request) if response is not None: - sys.stdout.write(json.dumps(response) + "\n") - sys.stdout.flush() - except KeyboardInterrupt: - break - except Exception as e: - logger.error(f"Server error: {e}") + print(json.dumps(response), flush=True) + except (json.JSONDecodeError, EOFError): + continue + except Exception: + logger.exception("MCP Loop error") if __name__ == "__main__": diff --git a/mempalace/normalize.py b/mempalace/normalize.py index 4252afa4dc..2b4691eb8a 100644 --- a/mempalace/normalize.py +++ b/mempalace/normalize.py @@ -41,9 +41,12 @@ "system-reminder", "command-message", "command-name", + "command-args", "task-notification", "user-prompt-submit-hook", "hook_output", + "local-command-caveat", + "local-command-stdout", ) @@ -89,6 +92,14 @@ def _tag_pattern(name: str) -> "re.Pattern[str]": # "… +N lines" collapsed-output marker, line-anchored. _COLLAPSED_LINES_RE = re.compile(r"(?m)^(?:> )?…\s*\+\d+ lines.*\n?") +# ANSI escape sequences leak into transcripts via Bash tool_result blocks +# (e.g. /context, /help, any colored CLI output). Each pattern is anchored on +# the literal ESC byte (\x1b) — user prose that mentions e.g. "[1m]" by name +# does not start with ESC and therefore stays intact. +_ANSI_CSI_RE = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]") +# OSC sequences (terminal title, hyperlinks) terminated by BEL or ST. +_ANSI_OSC_RE = re.compile(r"\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)") + def strip_noise(text: str) -> str: """Remove system tags, hook output, and Claude Code UI chrome from text. @@ -105,6 +116,10 @@ def strip_noise(text: str) -> str: # Strip the Claude Code collapsed-output chrome "[N tokens] (ctrl+o to expand)". # Narrow shape — a bare "(ctrl+o to expand)" in user prose stays intact. text = re.sub(r"\s*\[\d+\s+tokens?\]\s*\(ctrl\+o to expand\)", "", text) + # Strip ANSI escape sequences. Tag-wrapped ANSI is already gone via the + # tag patterns above; this sweep handles standalone ANSI in tool output. + text = _ANSI_CSI_RE.sub("", text) + text = _ANSI_OSC_RE.sub("", text) # Collapse runs of blank lines created by the removals text = re.sub(r"\n{4,}", "\n\n\n", text) return text.strip() diff --git a/mempalace/palace.py b/mempalace/palace.py index 07efb6a3e3..7cf4a36124 100644 --- a/mempalace/palace.py +++ b/mempalace/palace.py @@ -8,6 +8,8 @@ import hashlib import os import re +import threading +from typing import Optional from .backends.chroma import ChromaBackend @@ -52,10 +54,14 @@ def get_collection( palace_path: str, - collection_name: str = "mempalace_drawers", + collection_name: Optional[str] = None, create: bool = True, ): """Get the palace collection through the backend layer.""" + if collection_name is None: + from .config import get_configured_collection_name + + collection_name = get_configured_collection_name() return _DEFAULT_BACKEND.get_collection( palace_path, collection_name=collection_name, @@ -314,6 +320,47 @@ class MineAlreadyRunning(RuntimeError): """Raised when another `mempalace mine` already holds the per-palace lock.""" +# Per-thread record of palaces this thread already holds the lock for. Used by +# `mine_palace_lock` to short-circuit re-entrant acquisition from the same +# thread (e.g. miner.mine() acquires the outer lock then calls +# ChromaCollection.upsert which now also tries to acquire). Without this guard +# the inner call would block on its own outer flock (Linux fcntl locks are per +# open file description, so a same-thread second open of the lock file is a +# distinct lock and self-deadlocks). +# +# The holder set is tagged with ``pid`` so that a forked child does NOT +# inherit re-entrant credit from its parent: the OS-level flock IS NOT +# inherited as a "we hold it" semantically — the child must reacquire — but +# Python's ``threading.local`` IS inherited across fork. The pid check +# clears stale state so a forked child correctly hits the fcntl path. +_palace_lock_holders = threading.local() + + +def _holder_state(): + """Return the per-thread (pid, keys) record, refreshing after fork.""" + keys = getattr(_palace_lock_holders, "keys", None) + pid = getattr(_palace_lock_holders, "pid", None) + current_pid = os.getpid() + if keys is None or pid != current_pid: + keys = set() + _palace_lock_holders.keys = keys + _palace_lock_holders.pid = current_pid + return keys + + +def _held_by_this_thread(lock_key: str) -> bool: + """Return True if this thread already holds ``mine_palace_lock`` for ``lock_key``.""" + return lock_key in _holder_state() + + +def _mark_held(lock_key: str) -> None: + _holder_state().add(lock_key) + + +def _mark_released(lock_key: str) -> None: + _holder_state().discard(lock_key) + + @contextlib.contextmanager def mine_palace_lock(palace_path: str): """Per-palace non-blocking lock around the full `mine` pipeline. @@ -338,6 +385,12 @@ def mine_palace_lock(palace_path: str): Non-blocking: if another `mine` is already writing to this palace, raise MineAlreadyRunning so the caller can exit cleanly instead of piling up as a waiting worker. + + Re-entrant: if the current thread already holds the lock for the same + palace, the context manager passes through without re-acquiring. This + lets ChromaCollection write methods (which acquire the lock themselves + to protect MCP/direct callers) compose with miner.mine() (which holds + the outer lock for the entire mine pipeline) without self-deadlock. """ lock_dir = os.path.join(os.path.expanduser("~"), ".mempalace", "locks") os.makedirs(lock_dir, exist_ok=True) @@ -346,6 +399,11 @@ def mine_palace_lock(palace_path: str): palace_key = hashlib.sha256(lock_key_source.encode()).hexdigest()[:16] lock_path = os.path.join(lock_dir, f"mine_palace_{palace_key}.lock") + if _held_by_this_thread(palace_key): + # Same thread already holds the lock for this palace — pass through. + yield + return + lf = open(lock_path, "w") acquired = False try: @@ -369,7 +427,11 @@ def mine_palace_lock(palace_path: str): raise MineAlreadyRunning( f"another `mempalace mine` is already running against {resolved}" ) from exc - yield + _mark_held(palace_key) + try: + yield + finally: + _mark_released(palace_key) finally: if acquired: try: diff --git a/mempalace/repair.py b/mempalace/repair.py index 1cd1556975..10663b3acc 100644 --- a/mempalace/repair.py +++ b/mempalace/repair.py @@ -37,10 +37,13 @@ from datetime import datetime from typing import Optional +from chromadb.errors import NotFoundError as ChromaNotFoundError + from .backends.chroma import ChromaBackend, hnsw_capacity_status COLLECTION_NAME = "mempalace_drawers" +REPAIR_TEMP_COLLECTION = f"{COLLECTION_NAME}__repair_tmp" def _get_palace_path(): @@ -54,6 +57,16 @@ def _get_palace_path(): return default +def _get_collection_name() -> str: + """Resolve drawers collection name from config.""" + try: + from .config import get_configured_collection_name + + return get_configured_collection_name() + except Exception: + return COLLECTION_NAME + + def _paginate_ids(col, where=None): """Pull all IDs in a collection using pagination.""" ids = [] @@ -83,7 +96,111 @@ def _paginate_ids(col, where=None): return ids -def scan_palace(palace_path=None, only_wing=None): +def _extract_drawers(col, total: int, batch_size: int): + all_ids = [] + all_docs = [] + all_metas = [] + offset = 0 + while offset < total: + batch = col.get(limit=batch_size, offset=offset, include=["documents", "metadatas"]) + if not batch["ids"]: + break + all_ids.extend(batch["ids"]) + all_docs.extend(batch["documents"]) + all_metas.extend(batch["metadatas"]) + offset += len(batch["ids"]) + return all_ids, all_docs, all_metas + + +def _verify_collection_count(col, expected: int, label: str) -> None: + actual = col.count() + if actual != expected: + raise RuntimeError(f"{label} count mismatch: expected {expected}, got {actual}") + + +def _is_missing_collection_value_error(exc: ValueError) -> bool: + message = str(exc).lower() + return "does not exist" in message or "not found" in message + + +def _delete_collection_if_exists(backend, palace_path: str, collection_name: str) -> None: + try: + backend.delete_collection(palace_path, collection_name) + except ValueError as exc: + if _is_missing_collection_value_error(exc): + return + raise + except (FileNotFoundError, ChromaNotFoundError): + return + + +class RebuildCollectionError(RuntimeError): + """Raised when temp rebuild fails, carrying whether the live swap happened.""" + + def __init__(self, message: str, *, live_replaced: bool): + super().__init__(message) + self.live_replaced = live_replaced + + +def _rebuild_collection_via_temp( + backend, + palace_path: str, + all_ids, + all_docs, + all_metas, + batch_size: int, + collection_name: Optional[str] = None, + progress=print, +) -> int: + expected = len(all_ids) + collection_name = collection_name or _get_collection_name() + temp_name = f"{collection_name}__repair_tmp" + live_replaced = False + + try: + _delete_collection_if_exists(backend, palace_path, temp_name) + + progress(f" Building temporary collection: {temp_name}") + temp_col = backend.create_collection(palace_path, temp_name) + staged = 0 + for i in range(0, expected, batch_size): + batch_ids = all_ids[i : i + batch_size] + batch_docs = all_docs[i : i + batch_size] + batch_metas = all_metas[i : i + batch_size] + temp_col.upsert(documents=batch_docs, ids=batch_ids, metadatas=batch_metas) + staged += len(batch_ids) + progress(f" Staged {staged}/{expected} drawers...") + _verify_collection_count(temp_col, expected, "temporary rebuild") + + progress(" Rebuilding live collection...") + backend.delete_collection(palace_path, collection_name) + live_replaced = True + new_col = backend.create_collection(palace_path, collection_name) + + rebuilt = 0 + for i in range(0, expected, batch_size): + batch_ids = all_ids[i : i + batch_size] + batch_docs = all_docs[i : i + batch_size] + batch_metas = all_metas[i : i + batch_size] + new_col.upsert(documents=batch_docs, ids=batch_ids, metadatas=batch_metas) + rebuilt += len(batch_ids) + progress(f" Re-filed {rebuilt}/{expected} drawers...") + _verify_collection_count(new_col, expected, "rebuilt live collection") + + try: + _delete_collection_if_exists(backend, palace_path, temp_name) + except Exception: + pass + return rebuilt + except Exception as exc: + try: + _delete_collection_if_exists(backend, palace_path, temp_name) + except Exception: + pass + raise RebuildCollectionError(str(exc), live_replaced=live_replaced) from exc + + +def scan_palace(palace_path=None, only_wing=None, collection_name: Optional[str] = None): """Scan the palace for corrupt/unfetchable IDs. Probes in batches of 100, falls back to per-ID on failure. @@ -92,14 +209,15 @@ def scan_palace(palace_path=None, only_wing=None): Returns (good_set, bad_set). """ palace_path = palace_path or _get_palace_path() + collection_name = collection_name or _get_collection_name() print(f"\n Palace: {palace_path}") print(" Loading...") - col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME) + col = ChromaBackend().get_collection(palace_path, collection_name) where = {"wing": only_wing} if only_wing else None total = col.count() - print(f" Collection: {COLLECTION_NAME}, total: {total:,}") + print(f" Collection: {collection_name}, total: {total:,}") if only_wing: print(f" Scanning wing: {only_wing}") @@ -160,9 +278,10 @@ def scan_palace(palace_path=None, only_wing=None): return good_set, bad_set -def prune_corrupt(palace_path=None, confirm=False): +def prune_corrupt(palace_path=None, confirm=False, collection_name: Optional[str] = None): """Delete corrupt IDs listed in corrupt_ids.txt.""" palace_path = palace_path or _get_palace_path() + collection_name = collection_name or _get_collection_name() bad_file = os.path.join(palace_path, "corrupt_ids.txt") if not os.path.exists(bad_file): @@ -178,7 +297,7 @@ def prune_corrupt(palace_path=None, confirm=False): print(" Re-run with --confirm to actually delete.") return - col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME) + col = ChromaBackend().get_collection(palace_path, collection_name) before = col.count() print(f" Collection size before: {before:,}") @@ -232,7 +351,10 @@ def __init__(self, message: str, sqlite_count: "int | None", extracted: int): def check_extraction_safety( - palace_path: str, extracted: int, confirm_truncation_ok: bool = False + palace_path: str, + extracted: int, + confirm_truncation_ok: bool = False, + collection_name: Optional[str] = None, ) -> None: """Cross-check that ``extracted`` matches the SQLite ground truth. @@ -254,7 +376,8 @@ def check_extraction_safety( if confirm_truncation_ok: return - sqlite_count = sqlite_drawer_count(palace_path) + collection_name = collection_name or _get_collection_name() + sqlite_count = sqlite_drawer_count(palace_path, collection_name) cap_signal = extracted == CHROMADB_DEFAULT_GET_LIMIT if sqlite_count is not None and sqlite_count > extracted: @@ -290,7 +413,7 @@ def check_extraction_safety( raise TruncationDetected(message, sqlite_count, extracted) -def sqlite_drawer_count(palace_path: str) -> "int | None": +def sqlite_drawer_count(palace_path: str, collection_name: Optional[str] = None) -> "int | None": """Count rows in ``chroma.sqlite3.embeddings`` for the drawers collection. Used as an independent ground-truth check against the chromadb @@ -302,6 +425,7 @@ def sqlite_drawer_count(palace_path: str) -> "int | None": drift, missing tables, locked file). Callers treat ``None`` as "unknown" and fall back to the cap-detection check. """ + collection_name = collection_name or _get_collection_name() sqlite_path = os.path.join(palace_path, "chroma.sqlite3") if not os.path.exists(sqlite_path): return None @@ -318,7 +442,7 @@ def sqlite_drawer_count(palace_path: str) -> "int | None": JOIN collections c ON s.collection = c.id WHERE c.name = ? """, - (COLLECTION_NAME,), + (collection_name,), ).fetchone() return int(row[0]) if row and row[0] is not None else None finally: @@ -330,7 +454,11 @@ def sqlite_drawer_count(palace_path: str) -> "int | None": return None -def rebuild_index(palace_path=None, confirm_truncation_ok: bool = False): +def rebuild_index( + palace_path=None, + confirm_truncation_ok: bool = False, + collection_name: Optional[str] = None, +): """Rebuild the HNSW index from scratch. 1. Extract all drawers via ChromaDB get() @@ -345,6 +473,7 @@ def rebuild_index(palace_path=None, confirm_truncation_ok: bool = False): (typically only a concern for palaces sized at exactly 10 000 rows). """ palace_path = palace_path or _get_palace_path() + collection_name = collection_name or _get_collection_name() if not os.path.isdir(palace_path): print(f"\n No palace found at {palace_path}") @@ -357,7 +486,7 @@ def rebuild_index(palace_path=None, confirm_truncation_ok: bool = False): backend = ChromaBackend() try: - col = backend.get_collection(palace_path, COLLECTION_NAME) + col = backend.get_collection(palace_path, collection_name) total = col.count() except Exception as e: print(f" Error reading palace: {e}") @@ -373,18 +502,7 @@ def rebuild_index(palace_path=None, confirm_truncation_ok: bool = False): # Extract all drawers in batches print("\n Extracting drawers...") batch_size = 5000 - all_ids = [] - all_docs = [] - all_metas = [] - offset = 0 - while offset < total: - batch = col.get(limit=batch_size, offset=offset, include=["documents", "metadatas"]) - if not batch["ids"]: - break - all_ids.extend(batch["ids"]) - all_docs.extend(batch["documents"]) - all_metas.extend(batch["metadatas"]) - offset += len(batch["ids"]) + all_ids, all_docs, all_metas = _extract_drawers(col, total, batch_size) print(f" Extracted {len(all_ids)} drawers") # ── #1208 guard ────────────────────────────────────────────────── @@ -392,7 +510,12 @@ def rebuild_index(palace_path=None, confirm_truncation_ok: bool = False): # short of the SQLite ground truth (or when extraction == chromadb # default get() cap and the SQLite check couldn't run). try: - check_extraction_safety(palace_path, len(all_ids), confirm_truncation_ok) + check_extraction_safety( + palace_path, + len(all_ids), + confirm_truncation_ok, + collection_name=collection_name, + ) except TruncationDetected as e: print(e.message) return @@ -407,28 +530,34 @@ def rebuild_index(palace_path=None, confirm_truncation_ok: bool = False): # Rebuild with correct HNSW settings print(" Rebuilding collection with hnsw:space=cosine...") - backend.delete_collection(palace_path, COLLECTION_NAME) - new_col = backend.create_collection(palace_path, COLLECTION_NAME) - - filed = 0 try: - for i in range(0, len(all_ids), batch_size): - batch_ids = all_ids[i : i + batch_size] - batch_docs = all_docs[i : i + batch_size] - batch_metas = all_metas[i : i + batch_size] - new_col.upsert(documents=batch_docs, ids=batch_ids, metadatas=batch_metas) - filed += len(batch_ids) - print(f" Re-filed {filed}/{len(all_ids)} drawers...") - except Exception as e: + filed = _rebuild_collection_via_temp( + backend, + palace_path, + all_ids, + all_docs, + all_metas, + batch_size, + collection_name=collection_name, + progress=print, + ) + except RebuildCollectionError as e: print(f"\n ERROR during rebuild: {e}") - print(f" Only {filed}/{len(all_ids)} drawers were re-filed.") - if os.path.exists(backup_path): + print(" Rebuild aborted before completion.") + if e.live_replaced and os.path.exists(backup_path): print(f" Restoring from backup: {backup_path}") - backend.delete_collection(palace_path, COLLECTION_NAME) - shutil.copy2(backup_path, sqlite_path) - print(" Backup restored. Palace is back to pre-repair state.") - else: + try: + _close_chroma_handles(palace_path, backend=backend) + _delete_collection_if_exists(backend, palace_path, collection_name) + shutil.copy2(backup_path, sqlite_path) + print(" Backup restored. Palace is back to pre-repair state.") + except Exception as restore_error: + print(f" Backup restore failed: {restore_error}") + print(f" Manual restore required from: {backup_path}") + elif e.live_replaced: print(" No backup available. Re-mine from source files to recover.") + else: + print(" Live collection was not replaced; leaving the original palace untouched.") raise print(f"\n Repair complete. {filed} drawers rebuilt.") @@ -436,7 +565,7 @@ def rebuild_index(palace_path=None, confirm_truncation_ok: bool = False): print(f"\n{'=' * 55}\n") -def status(palace_path=None) -> dict: +def status(palace_path=None, collection_name: Optional[str] = None) -> dict: """Read-only health check: compare sqlite vs HNSW element counts. Catches the #1222 failure mode where chromadb's HNSW segment freezes @@ -454,6 +583,7 @@ def status(palace_path=None) -> dict: ``status="unknown"`` when no palace exists at the given path. """ palace_path = palace_path or _get_palace_path() + collection_name = collection_name or _get_collection_name() print(f"\n{'=' * 55}") print(" MemPalace Repair — Status") print(f"{'=' * 55}\n") @@ -463,7 +593,7 @@ def status(palace_path=None) -> dict: print(" No palace found.\n") return {"status": "unknown", "message": "no palace at path"} - drawers = hnsw_capacity_status(palace_path, "mempalace_drawers") + drawers = hnsw_capacity_status(palace_path, collection_name) closets = hnsw_capacity_status(palace_path, "mempalace_closets") for label, info in (("drawers", drawers), ("closets", closets)): @@ -494,12 +624,18 @@ def status(palace_path=None) -> dict: # --------------------------------------------------------------------------- -def _close_chroma_handles(palace_path: str) -> None: - """Drop ChromaBackend + chromadb singleton caches so OS mmap handles release.""" +def _close_chroma_handles(palace_path: str, backend: ChromaBackend | None = None) -> None: + """Drop ChromaBackend + chromadb singleton caches so OS mmap handles release. + + When ``backend`` is provided, close the live instance so rollback/restore + releases the handles it was already using. Otherwise fall back to a + transient backend instance for the max-seq-id repair path. + """ import gc try: - ChromaBackend().close_palace(palace_path) + closer = backend if backend is not None else ChromaBackend() + closer.close_palace(palace_path) except Exception: pass try: @@ -546,7 +682,10 @@ def _detect_poisoned_max_seq_ids( "SELECT segment_id, seq_id FROM max_seq_id WHERE seq_id > ?", (threshold,), ).fetchall() - return [(str(sid), int(val)) for sid, val in rows] + return [ + (str(sid), (int.from_bytes(val, "big") if isinstance(val, (bytes, bytearray)) else int(val))) + for sid, val in rows + ] def _compute_heuristic_seq_id(cur: sqlite3.Cursor, segment_id: str) -> int: diff --git a/mempalace/searcher.py b/mempalace/searcher.py index c08a6ce6bc..7aa909575e 100644 --- a/mempalace/searcher.py +++ b/mempalace/searcher.py @@ -372,6 +372,7 @@ def _bm25_only_via_sqlite( room: str = None, n_results: int = 5, max_candidates: int = 500, + collection_name: str = None, ) -> dict: """BM25-only search reading drawers directly from chroma.sqlite3. @@ -395,6 +396,35 @@ def _bm25_only_via_sqlite( "error": "No palace found", "hint": "Run: mempalace init && mempalace mine ", } + if collection_name is None: + from .config import get_configured_collection_name + + collection_name = get_configured_collection_name() + + def _metadata_filter_sql(row_id_expr: str) -> tuple[str, list[str]]: + clauses = [] + params = [] + for key, value in (("wing", wing), ("room", room)): + if not value: + continue + clauses.append( + f""" + AND EXISTS ( + SELECT 1 + FROM embedding_metadata mf + WHERE mf.id = {row_id_expr} + AND mf.key = ? + AND COALESCE( + mf.string_value, + CAST(mf.int_value AS TEXT), + CAST(mf.float_value AS TEXT), + CAST(mf.bool_value AS TEXT) + ) = ? + ) + """ + ) + params.extend([key, value]) + return "".join(clauses), params try: conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) @@ -406,45 +436,57 @@ def _bm25_only_via_sqlite( # shorter than 3 chars (trigram tokenizer can't match them). tokens = [t for t in _tokenize(query) if len(t) >= 3] candidate_ids: list[int] = [] + use_recency_fallback = not tokens if tokens: fts_query = " OR ".join(tokens) + filter_sql, filter_params = _metadata_filter_sql("embedding_fulltext_search.rowid") try: rows = conn.execute( - """ - SELECT rowid + f""" + SELECT embedding_fulltext_search.rowid FROM embedding_fulltext_search + JOIN embeddings e ON e.id = embedding_fulltext_search.rowid + JOIN segments s ON e.segment_id = s.id + JOIN collections c ON s.collection = c.id WHERE embedding_fulltext_search MATCH ? + AND c.name = ? + {filter_sql} LIMIT ? """, - (fts_query, max_candidates), + (fts_query, collection_name, *filter_params, max_candidates), ).fetchall() candidate_ids = [r[0] for r in rows] except sqlite3.Error: # FTS5 tokenizer mismatch or syntax error — fall through # to the recency-window selector below. logger.debug("FTS5 MATCH failed; using recency fallback", exc_info=True) - - if not candidate_ids: - # No FTS hits (or no usable tokens) — pull the most recent - # rows for the drawers segment so we can BM25-rank something - # rather than return empty-handed. Wrapped in try/except - # because the schema may differ on legacy palaces (older - # chromadb without ``created_at``, missing ``segments`` - # rows after partial restore, etc.); on schema mismatch we - # fall back to ordering by primary-key id and finally to an - # empty result rather than letting search raise. + use_recency_fallback = True + + if not candidate_ids and use_recency_fallback: + # No usable FTS tokens, or FTS itself failed — pull the most + # recent rows for the drawers segment so we can BM25-rank + # something rather than return empty-handed. A clean FTS miss + # must stay empty, especially after wing/room filtering, because + # recency fallback would return unrelated scoped drawers. + # Wrapped in try/except because the schema may differ on legacy + # palaces (older chromadb without ``created_at``, missing + # ``segments`` rows after partial restore, etc.); on schema + # mismatch we fall back to ordering by primary-key id and finally + # to an empty result rather than letting search raise. try: + filter_sql, filter_params = _metadata_filter_sql("e.id") rows = conn.execute( - """ + f""" SELECT e.id FROM embeddings e JOIN segments s ON e.segment_id = s.id JOIN collections c ON s.collection = c.id - WHERE c.name = 'mempalace_drawers' + WHERE c.name = ? + {filter_sql} ORDER BY e.created_at DESC LIMIT ? """, - (max_candidates,), + (collection_name, *filter_params, max_candidates), ).fetchall() candidate_ids = [r[0] for r in rows] except sqlite3.Error: @@ -453,17 +495,19 @@ def _bm25_only_via_sqlite( exc_info=True, ) try: + filter_sql, filter_params = _metadata_filter_sql("e.id") rows = conn.execute( - """ + f""" SELECT e.id FROM embeddings e JOIN segments s ON e.segment_id = s.id JOIN collections c ON s.collection = c.id - WHERE c.name = 'mempalace_drawers' + WHERE c.name = ? + {filter_sql} ORDER BY e.id DESC LIMIT ? """, - (max_candidates,), + (collection_name, *filter_params, max_candidates), ).fetchall() candidate_ids = [r[0] for r in rows] except sqlite3.Error: @@ -553,6 +597,7 @@ def search_memories( n_results: int = 5, max_distance: float = 0.0, vector_disabled: bool = False, + collection_name: str = None, ) -> dict: """Programmatic search — returns a dict instead of printing. @@ -580,10 +625,11 @@ def search_memories( wing=wing, room=room, n_results=n_results, + collection_name=collection_name, ) try: - drawers_col = get_collection(palace_path, create=False) + drawers_col = get_collection(palace_path, collection_name=collection_name, create=False) except Exception as e: logger.error("No palace found at %s: %s", palace_path, e) return { diff --git a/tests/test_backends.py b/tests/test_backends.py index 5efa71b319..d94fecb72e 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -1,4 +1,5 @@ import os +import pickle import sqlite3 from pathlib import Path @@ -18,6 +19,7 @@ ChromaCollection, _fix_blob_seq_ids, _pin_hnsw_threads, + quarantine_invalid_hnsw_metadata, quarantine_stale_hnsw, ) @@ -708,7 +710,10 @@ def test_make_client_quarantines_only_on_first_call_per_palace(tmp_path, monkeyp """Quarantine fires on first ``make_client()`` for a palace, then is skipped on subsequent calls — prevents runtime thrash where a daemon's own steady writes bump ``chroma.sqlite3`` faster than HNSW flushes, - making the mtime heuristic falsely trigger every reconnect.""" + making the mtime heuristic falsely trigger every reconnect. + + Invalid metadata quarantine shares the same cold-start gate here; the + more aggressive refresh path lives in ``_client()``.""" from mempalace.backends.chroma import ChromaBackend palace_path = str(tmp_path / "palace") @@ -730,9 +735,37 @@ def _spy(path, stale_seconds=300.0): ChromaBackend.make_client(palace_path) ChromaBackend.make_client(palace_path) - assert calls == [ - palace_path - ], "quarantine_stale_hnsw should fire once per palace per process, not on every reconnect" + assert calls == [palace_path], ( + "quarantine_stale_hnsw should fire once per palace per process, not on every reconnect" + ) + + +def test_make_client_gates_invalid_metadata_on_first_call(tmp_path, monkeypatch): + """Invalid metadata quarantine is gated on the first make_client() call.""" + from mempalace.backends.chroma import ChromaBackend + + palace_path = str(tmp_path / "palace") + os.makedirs(palace_path, exist_ok=True) + (Path(palace_path) / "chroma.sqlite3").write_text("") + + monkeypatch.setattr(ChromaBackend, "_quarantined_paths", set()) + + calls: list[str] = [] + + def _invalid(path, *args, **kwargs): + calls.append(path) + return [] + + def _stale(path, stale_seconds=300.0): + return [] + + monkeypatch.setattr("mempalace.backends.chroma.quarantine_invalid_hnsw_metadata", _invalid) + monkeypatch.setattr("mempalace.backends.chroma.quarantine_stale_hnsw", _stale) + + ChromaBackend.make_client(palace_path) + ChromaBackend.make_client(palace_path) + + assert calls == [palace_path] def test_make_client_quarantines_each_palace_independently(tmp_path, monkeypatch): @@ -811,3 +844,268 @@ def test_get_collection_applies_retrofit_on_existing_palace(tmp_path): ) assert wrapper._collection.configuration_json["hnsw"]["num_threads"] == 1 + + +def test_quarantine_invalid_hnsw_metadata_renames_missing_dimensionality(tmp_path): + palace = tmp_path / "palace" + palace.mkdir() + seg = palace / "abcd-1234-5678" + seg.mkdir() + with open(seg / "index_metadata.pickle", "wb") as f: + pickle.dump({"dimensionality": None, "id_to_label": {"a": 1}}, f) + + moved = quarantine_invalid_hnsw_metadata(str(palace)) + + assert len(moved) == 1 + assert ".corrupt-" in moved[0] + assert not seg.exists() + + +def test_quarantine_invalid_hnsw_metadata_allows_uninitialized_segment(tmp_path): + palace = tmp_path / "palace" + palace.mkdir() + seg = palace / "abcd-1234-5678" + seg.mkdir() + with open(seg / "index_metadata.pickle", "wb") as f: + pickle.dump({"dimensionality": None, "id_to_label": {}}, f) + + moved = quarantine_invalid_hnsw_metadata(str(palace)) + + assert moved == [] + assert seg.exists() + + +def test_quarantine_invalid_hnsw_metadata_rejects_non_dict_id_to_label(tmp_path): + palace = tmp_path / "palace" + palace.mkdir() + seg = palace / "abcd-1234-5678" + seg.mkdir() + with open(seg / "index_metadata.pickle", "wb") as f: + pickle.dump({"dimensionality": 8, "id_to_label": ["a", "b"]}, f) + + moved = quarantine_invalid_hnsw_metadata(str(palace)) + + assert len(moved) == 1 + assert ".corrupt-" in moved[0] + assert not seg.exists() + + +def test_quarantine_invalid_hnsw_metadata_rejects_non_schema_payload(tmp_path): + palace = tmp_path / "palace" + palace.mkdir() + seg = palace / "abcd-1234-5678" + seg.mkdir() + with open(seg / "index_metadata.pickle", "wb") as f: + pickle.dump(["not", "a", "metadata", "object"], f) + + moved = quarantine_invalid_hnsw_metadata(str(palace)) + + assert len(moved) == 1 + assert ".corrupt-" in moved[0] + assert not seg.exists() + + +def _dangerous_pickle_payload_executed(): + raise AssertionError("unsafe pickle payload executed") + + +class _DangerousPickle: + def __reduce__(self): + return (_dangerous_pickle_payload_executed, ()) + + +def test_quarantine_invalid_hnsw_metadata_rejects_unsafe_pickle(tmp_path): + palace = tmp_path / "palace" + palace.mkdir() + seg = palace / "abcd-1234-5678" + seg.mkdir() + with open(seg / "index_metadata.pickle", "wb") as f: + pickle.dump(_DangerousPickle(), f) + + moved = quarantine_invalid_hnsw_metadata(str(palace)) + + assert len(moved) == 1 + assert ".corrupt-" in moved[0] + assert not seg.exists() + + +def test_quarantine_invalid_hnsw_metadata_skips_transient_read_errors(tmp_path, monkeypatch): + palace = tmp_path / "palace" + palace.mkdir() + seg = palace / "abcd-1234-5678" + seg.mkdir() + meta = seg / "index_metadata.pickle" + meta.write_bytes(b"partial") + + monkeypatch.setattr( + "mempalace.backends.chroma._SafePersistentDataUnpickler.load", + lambda path: (_ for _ in ()).throw(EOFError("flush in progress")), + ) + + moved = quarantine_invalid_hnsw_metadata(str(palace)) + + assert moved == [] + assert seg.exists() + + +def test_quarantine_invalid_hnsw_metadata_skips_truncated_pickle(tmp_path, monkeypatch): + palace = tmp_path / "palace" + palace.mkdir() + seg = palace / "abcd-1234-5678" + seg.mkdir() + meta = seg / "index_metadata.pickle" + meta.write_bytes(b"partial") + + monkeypatch.setattr( + "mempalace.backends.chroma._SafePersistentDataUnpickler.load", + lambda path: (_ for _ in ()).throw(pickle.UnpicklingError("pickle data was truncated")), + ) + + moved = quarantine_invalid_hnsw_metadata(str(palace)) + + assert moved == [] + assert seg.exists() + + +def test_chroma_backend_preflights_metadata_before_persistent_client(tmp_path, monkeypatch): + palace = tmp_path / "palace" + palace.mkdir() + calls = [] + + def _record(name): + def inner(path, *args, **kwargs): + calls.append((name, path)) + return [] if name != "blob" else None + + return inner + + monkeypatch.setattr("mempalace.backends.chroma._fix_blob_seq_ids", _record("blob")) + monkeypatch.setattr( + "mempalace.backends.chroma.quarantine_invalid_hnsw_metadata", _record("invalid") + ) + monkeypatch.setattr("mempalace.backends.chroma.quarantine_stale_hnsw", _record("stale")) + + class DummyClient: + pass + + monkeypatch.setattr( + "mempalace.backends.chroma.chromadb.PersistentClient", lambda path: DummyClient() + ) + + backend = ChromaBackend() + backend._client(str(palace)) + + assert calls == [ + ("blob", str(palace)), + ("invalid", str(palace)), + ("stale", str(palace)), + ] + + +def test_chroma_backend_stale_quarantine_is_cold_start_only_on_refresh(tmp_path, monkeypatch): + palace = tmp_path / "palace" + palace.mkdir() + (palace / "chroma.sqlite3").write_text("") + calls = [] + + def _record(name): + def inner(path, *args, **kwargs): + calls.append((name, path)) + return [] if name != "blob" else None + + return inner + + monkeypatch.setattr(ChromaBackend, "_quarantined_paths", set()) + monkeypatch.setattr("mempalace.backends.chroma._fix_blob_seq_ids", _record("blob")) + monkeypatch.setattr( + "mempalace.backends.chroma.quarantine_invalid_hnsw_metadata", _record("invalid") + ) + monkeypatch.setattr("mempalace.backends.chroma.quarantine_stale_hnsw", _record("stale")) + + class DummyClient: + pass + + monkeypatch.setattr( + "mempalace.backends.chroma.chromadb.PersistentClient", lambda path: DummyClient() + ) + + backend = ChromaBackend() + stats = iter([(1, 1.0), (1, 1.0), (1, 2.0), (1, 2.0)]) + monkeypatch.setattr(backend, "_db_stat", lambda path: next(stats)) + + backend._client(str(palace)) + backend._client(str(palace)) + + assert calls == [ + ("blob", str(palace)), + ("invalid", str(palace)), + ("stale", str(palace)), + ("blob", str(palace)), + ] + + +def test_chroma_backend_requarantines_after_inode_replacement(tmp_path, monkeypatch): + palace = tmp_path / "palace" + palace.mkdir() + (palace / "chroma.sqlite3").write_text("") + calls = [] + + def _record(name): + def inner(path, *args, **kwargs): + calls.append((name, path)) + return [] if name != "blob" else None + + return inner + + monkeypatch.setattr(ChromaBackend, "_quarantined_paths", set()) + monkeypatch.setattr("mempalace.backends.chroma._fix_blob_seq_ids", _record("blob")) + monkeypatch.setattr( + "mempalace.backends.chroma.quarantine_invalid_hnsw_metadata", _record("invalid") + ) + monkeypatch.setattr("mempalace.backends.chroma.quarantine_stale_hnsw", _record("stale")) + + class DummyClient: + pass + + monkeypatch.setattr( + "mempalace.backends.chroma.chromadb.PersistentClient", lambda path: DummyClient() + ) + + backend = ChromaBackend() + stats = iter([(1, 1.0), (1, 1.0), (2, 2.0), (2, 2.0)]) + monkeypatch.setattr(backend, "_db_stat", lambda path: next(stats)) + + backend._client(str(palace)) + backend._client(str(palace)) + + assert calls == [ + ("blob", str(palace)), + ("invalid", str(palace)), + ("stale", str(palace)), + ("blob", str(palace)), + ("invalid", str(palace)), + ("stale", str(palace)), + ] + + +def test_palace_get_collection_uses_configured_collection_name(monkeypatch): + from mempalace import palace + + captured = {} + + def fake_get_collection(palace_path, collection_name=None, create=False): + captured["palace_path"] = palace_path + captured["collection_name"] = collection_name + captured["create"] = create + return object() + + monkeypatch.setattr(palace._DEFAULT_BACKEND, "get_collection", fake_get_collection) + monkeypatch.setattr("mempalace.config.get_configured_collection_name", lambda: "custom_drawers") + + palace.get_collection("/palace", create=False) + + assert captured == { + "palace_path": "/palace", + "collection_name": "custom_drawers", + "create": False, + } diff --git a/tests/test_chroma_collection_lock.py b/tests/test_chroma_collection_lock.py new file mode 100644 index 0000000000..b5d30fbcfe --- /dev/null +++ b/tests/test_chroma_collection_lock.py @@ -0,0 +1,327 @@ +"""Tests for ChromaCollection's palace-write-lock integration. + +Closes the gap left by ``mine_palace_lock`` only protecting the +``mempalace mine`` pipeline: MCP/direct writers that call +``ChromaCollection.add/upsert/update/delete`` must also serialize against +mine and against each other to avoid the multi-threaded HNSW corruption +documented in #974/#965. + +Property tested: + +* ``ChromaCollection(c, palace_path=p)`` wraps every write with + ``mine_palace_lock(p)``. +* Writes raise ``MineAlreadyRunning`` when another holder owns the lock + (instead of silently racing into the underlying chromadb call). +* Re-entrant composition with ``miner.mine()`` does not self-deadlock: + ``with mine_palace_lock(p): col.upsert(...)`` runs to completion. +* ``ChromaCollection(c)`` (no palace_path) preserves legacy no-lock + behaviour for tests/callers that build the adapter directly without + going through ``ChromaBackend``. + +POSIX-only: ``mine_palace_lock`` uses ``fcntl`` on Unix and ``msvcrt`` on +Windows; the contention semantics differ enough that the cross-process +tests are skipped on Windows runners. +""" + +from __future__ import annotations + +import multiprocessing +import os +import time + +import pytest + +from mempalace.backends.chroma import ChromaCollection +from mempalace.palace import MineAlreadyRunning, mine_palace_lock + + +def _get_mp_context(): + """Same start-method picker as test_palace_locks.py.""" + start_method = "spawn" if os.name == "nt" else "fork" + return multiprocessing.get_context(start_method) + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class _FakeChromaCollection: + """Records calls; never blocks. Stand-in for chromadb.Collection.""" + + def __init__(self): + self.adds: list[dict] = [] + self.upserts: list[dict] = [] + self.updates: list[dict] = [] + self.deletes: list[dict] = [] + + def add(self, **kwargs): + self.adds.append(kwargs) + + def upsert(self, **kwargs): + self.upserts.append(kwargs) + + def update(self, **kwargs): + self.updates.append(kwargs) + + def delete(self, **kwargs): + self.deletes.append(kwargs) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _hold_lock(palace_path: str, ready_flag: str, release_flag: str) -> int: + """Acquire ``mine_palace_lock``, signal readiness, wait for release. + + Mirrors the helper in ``test_palace_locks.py`` so the contention + semantics match across both test files. + """ + try: + with mine_palace_lock(palace_path): + open(ready_flag, "w").close() + for _ in range(500): + if os.path.exists(release_flag): + return 0 + time.sleep(0.01) + return 0 + except MineAlreadyRunning: + return 1 + + +# --------------------------------------------------------------------------- +# Tests — opt-in lock wiring +# --------------------------------------------------------------------------- + + +def test_palace_path_none_skips_lock(tmp_path, monkeypatch): + """Legacy callers (``ChromaCollection(c)``) keep no-lock behaviour. + + A ``ChromaCollection`` built without ``palace_path`` must not touch the + lock infrastructure at all. This guards against regressions where a + test or third-party caller relies on the historical bare-write path. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + fake = _FakeChromaCollection() + col = ChromaCollection(fake) # no palace_path -> no lock + + # Hold the lock in a child process. Without palace_path, the parent + # write must still succeed (the lock does not gate this caller). + palace = str(tmp_path / "palace") + ready = str(tmp_path / "ready") + release = str(tmp_path / "release") + ctx = _get_mp_context() + holder = ctx.Process(target=_hold_lock, args=(palace, ready, release)) + holder.start() + try: + for _ in range(500): + if os.path.exists(ready): + break + time.sleep(0.01) + assert os.path.exists(ready), "holder failed to acquire lock" + + col.upsert(documents=["doc"], ids=["id-1"]) + assert fake.upserts == [{"documents": ["doc"], "ids": ["id-1"]}] + finally: + open(release, "w").close() + holder.join(timeout=5) + + +def test_writer_blocks_during_mine(tmp_path, monkeypatch): + """A held ``mine_palace_lock`` causes ``ChromaCollection`` writes to raise. + + This is the property that closes the MCP-bypass gap: when a mine is in + flight, MCP/direct writes raise ``MineAlreadyRunning`` rather than + silently entering chromadb's write path concurrent with mine. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + palace = str(tmp_path / "palace") + ready = str(tmp_path / "ready") + release = str(tmp_path / "release") + + ctx = _get_mp_context() + holder = ctx.Process(target=_hold_lock, args=(palace, ready, release)) + holder.start() + try: + for _ in range(500): + if os.path.exists(ready): + break + time.sleep(0.01) + assert os.path.exists(ready), "holder failed to acquire lock" + + fake = _FakeChromaCollection() + col = ChromaCollection(fake, palace_path=palace) + + with pytest.raises(MineAlreadyRunning): + col.upsert(documents=["doc"], ids=["id-1"]) + with pytest.raises(MineAlreadyRunning): + col.add(documents=["doc"], ids=["id-2"]) + with pytest.raises(MineAlreadyRunning): + col.update(ids=["id-3"], documents=["doc"]) + with pytest.raises(MineAlreadyRunning): + col.delete(ids=["id-4"]) + + # The fake must have received NO calls — the lock must gate + # before reaching the underlying chromadb layer. + assert fake.upserts == [] + assert fake.adds == [] + assert fake.updates == [] + assert fake.deletes == [] + finally: + open(release, "w").close() + holder.join(timeout=5) + + +def test_reentrant_inside_mine_passes_through(tmp_path, monkeypatch): + """``ChromaCollection.upsert`` inside ``mine_palace_lock`` does not deadlock. + + ``miner.mine()`` already holds ``mine_palace_lock(palace_path)`` for the + full mine pipeline; ``_mine_body`` then calls + ``collection.upsert(...)``. With the per-thread re-entrant guard in + ``mine_palace_lock``, the inner acquire is a pass-through and the + underlying chromadb call runs immediately. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + palace = str(tmp_path / "palace") + fake = _FakeChromaCollection() + col = ChromaCollection(fake, palace_path=palace) + + with mine_palace_lock(palace): + # If the re-entrant guard were missing, this would self-deadlock on + # the underlying flock. We rely on pytest-timeout (configured in + # pyproject.toml) to enforce this in CI; the assertion just confirms + # the call landed. + col.upsert(documents=["d"], ids=["i"], metadatas=[{"k": "v"}]) + col.add(documents=["d2"], ids=["i2"]) + col.update(ids=["i"], documents=["d-updated"]) + col.delete(ids=["i2"]) + + assert len(fake.upserts) == 1 + assert len(fake.adds) == 1 + assert len(fake.updates) == 1 + assert len(fake.deletes) == 1 + + +class _SlowFakeChromaCollection(_FakeChromaCollection): + """Fake whose write methods hold the caller for ``hold_seconds``. + + Used to keep ``mine_palace_lock`` acquired long enough for a sibling + process to contend deterministically. + """ + + def __init__(self, hold_seconds: float = 0.3): + super().__init__() + self._hold = hold_seconds + + def upsert(self, **kwargs): + time.sleep(self._hold) + super().upsert(**kwargs) + + +def _slow_writer_target(palace_path, tmp_path_str, pid, result_q): + """Subprocess target: try a slow upsert, report ok/busy.""" + os.environ["HOME"] = tmp_path_str + # Fresh import inside child so HOME monkeypatch routes the lock dir. + from mempalace.backends.chroma import ChromaCollection as _CC + from mempalace.palace import MineAlreadyRunning as _MAR + + fake = _SlowFakeChromaCollection(hold_seconds=0.3) + col = _CC(fake, palace_path=palace_path) + try: + col.upsert(documents=[f"d{pid}"], ids=[f"i{pid}"]) + result_q.put(("ok", pid)) + except _MAR: + result_q.put(("busy", pid)) + + +def test_concurrent_writers_serialize(tmp_path, monkeypatch): + """Two processes calling ``ChromaCollection.upsert`` against the same + palace must be serialized: at most one enters chromadb at a time, the + other raises ``MineAlreadyRunning``. + + This is the property that prevents the parallel HNSW insert race that + drives #974/#965 — under concurrent MCP write fan-out, exactly one + writer reaches chromadb and the rest fail loudly instead of corrupting + the index. + + The slow fake holds the lock for 0.3s per writer, large enough for the + second process to contend even on slow CI runners. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + palace = str(tmp_path / "palace") + + ctx = _get_mp_context() + result_q = ctx.Queue() + + p1 = ctx.Process( + target=_slow_writer_target, args=(palace, str(tmp_path), 1, result_q) + ) + p2 = ctx.Process( + target=_slow_writer_target, args=(palace, str(tmp_path), 2, result_q) + ) + p1.start() + # Tiny stagger so p1 wins the race deterministically; without it the + # OS scheduler can pick either, which is also a valid outcome but + # makes the assertion brittle on slow CI. + time.sleep(0.05) + p2.start() + p1.join(timeout=5) + p2.join(timeout=5) + + outcomes = [result_q.get(timeout=1) for _ in range(2)] + statuses = sorted(o[0] for o in outcomes) + assert statuses == ["busy", "ok"], ( + f"expected one ok + one busy, got {outcomes}" + ) + + +def test_read_path_does_not_acquire_lock(tmp_path, monkeypatch): + """``query`` / ``get`` / ``count`` must not be gated by the write lock. + + Read traffic is the dominant workload (semantic search, MCP get, etc.) + and serializing it against mine would tank latency for no correctness + benefit. This test pins that property: with another process holding + the write lock, reads must still complete instantly. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + palace = str(tmp_path / "palace") + ready = str(tmp_path / "ready") + release = str(tmp_path / "release") + + ctx = _get_mp_context() + holder = ctx.Process(target=_hold_lock, args=(palace, ready, release)) + holder.start() + try: + for _ in range(500): + if os.path.exists(ready): + break + time.sleep(0.01) + assert os.path.exists(ready), "holder failed to acquire lock" + + # _FakeChromaCollection doesn't implement query/get/count; we only + # need to confirm the wrapper does not call into mine_palace_lock + # for reads, which we assert by observing the wrapped methods are + # NOT in ChromaCollection's _write_lock path. A direct check via + # source inspection is more honest than mocking the entire chroma + # surface here. + import inspect + + from mempalace.backends.chroma import ChromaCollection as _CC + + for write_attr in ("add", "upsert", "update", "delete"): + src = inspect.getsource(getattr(_CC, write_attr)) + assert "_write_lock" in src, f"{write_attr} should acquire write lock" + + for read_attr in ("query", "get", "count"): + method = getattr(_CC, read_attr, None) + if method is None: + continue + src = inspect.getsource(method) + assert "_write_lock" not in src, ( + f"{read_attr} must NOT acquire the write lock (read path)" + ) + finally: + open(release, "w").close() + holder.join(timeout=5) diff --git a/tests/test_cli.py b/tests/test_cli.py index af7b39d0a6..f61b4b95a7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4,7 +4,7 @@ import shlex import sys from pathlib import Path -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest @@ -463,6 +463,7 @@ def test_cmd_mine_convos_mode(mock_config_cls): no_gitignore=False, include_ignored=[], extract="general", + include_subagents=False, ) with patch("mempalace.convo_miner.mine_convos") as mock_mine: cmd_mine(args) @@ -474,9 +475,32 @@ def test_cmd_mine_convos_mode(mock_config_cls): limit=10, dry_run=True, extract_mode="general", + include_subagents=False, ) +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_mine_convos_mode_threads_include_subagents_flag(mock_config_cls): + mock_config_cls.return_value.palace_path = "/fake/palace" + args = argparse.Namespace( + dir="/chats", + palace=None, + mode="convos", + wing="mywing", + agent="me", + limit=10, + dry_run=True, + no_gitignore=False, + include_ignored=[], + extract="exchange", + include_subagents=True, + ) + with patch("mempalace.convo_miner.mine_convos") as mock_mine: + cmd_mine(args) + kwargs = mock_mine.call_args.kwargs + assert kwargs["include_subagents"] is True + + @patch("mempalace.cli.MempalaceConfig") def test_cmd_mine_include_ignored_comma_split(mock_config_cls): mock_config_cls.return_value.palace_path = "/fake/palace" @@ -721,6 +745,7 @@ def test_cmd_repair_error_reading(mock_config_cls, tmp_path, capsys): palace_dir.mkdir() (palace_dir / "chroma.sqlite3").write_text("db") mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" args = argparse.Namespace(palace=None) mock_backend = MagicMock() mock_backend.get_collection.side_effect = Exception("corrupt db") @@ -736,6 +761,7 @@ def test_cmd_repair_zero_drawers(mock_config_cls, tmp_path, capsys): palace_dir.mkdir() (palace_dir / "chroma.sqlite3").write_text("db") mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" args = argparse.Namespace(palace=None) mock_col = MagicMock() mock_col.count.return_value = 0 @@ -752,6 +778,7 @@ def test_cmd_repair_success(mock_config_cls, tmp_path, capsys): palace_dir.mkdir() (palace_dir / "chroma.sqlite3").write_text("db") mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" args = argparse.Namespace(palace=None, yes=True) mock_col = MagicMock() mock_col.count.return_value = 2 @@ -760,13 +787,98 @@ def test_cmd_repair_success(mock_config_cls, tmp_path, capsys): "documents": ["doc1", "doc2"], "metadatas": [{"wing": "a"}, {"wing": "b"}], } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 mock_new_col = MagicMock() + mock_new_col.count.return_value = 2 mock_backend = _mock_backend_for(col=mock_col, new_col=mock_new_col) + mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] with patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend): cmd_repair(args) out = capsys.readouterr().out assert "Repair complete" in out assert "2 drawers rebuilt" in out + assert mock_backend.delete_collection.call_args_list == [ + call(str(palace_dir), "mempalace_drawers__repair_tmp"), + call(str(palace_dir), "mempalace_drawers"), + call(str(palace_dir), "mempalace_drawers__repair_tmp"), + ] + mock_temp_col.upsert.assert_called_once() + mock_new_col.upsert.assert_called_once() + mock_new_col.add.assert_not_called() + + +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_repair_uses_configured_collection(mock_config_cls, tmp_path, capsys): + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + (palace_dir / "chroma.sqlite3").write_text("db") + mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "custom_drawers" + args = argparse.Namespace(palace=None, yes=True) + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_new_col = MagicMock() + mock_new_col.count.return_value = 2 + mock_backend = _mock_backend_for(col=mock_col, new_col=mock_new_col) + mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] + + with patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend): + cmd_repair(args) + + out = capsys.readouterr().out + assert "Repair complete" in out + mock_backend.get_collection.assert_called_once_with(str(palace_dir), "custom_drawers") + assert mock_backend.create_collection.call_args_list == [ + call(str(palace_dir), "custom_drawers__repair_tmp"), + call(str(palace_dir), "custom_drawers"), + ] + assert mock_backend.delete_collection.call_args_list == [ + call(str(palace_dir), "custom_drawers__repair_tmp"), + call(str(palace_dir), "custom_drawers"), + call(str(palace_dir), "custom_drawers__repair_tmp"), + ] + + +@patch("mempalace.cli.MempalaceConfig") +def test_cmd_repair_restores_backup_on_live_rebuild_failure(mock_config_cls, tmp_path, capsys): + palace_dir = tmp_path / "palace" + palace_dir.mkdir() + (palace_dir / "chroma.sqlite3").write_text("db") + mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" + args = argparse.Namespace(palace=None, yes=True) + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_backend = _mock_backend_for(col=mock_col) + mock_backend.create_collection.side_effect = [mock_temp_col, RuntimeError("live build failed")] + with patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend): + with pytest.raises(SystemExit) as excinfo: + cmd_repair(args) + out = capsys.readouterr().out + assert excinfo.value.code == 1 + assert "Repair failed" in out + assert "restoring from backup" in out + mock_backend.close_palace.assert_called_once_with(str(palace_dir)) + assert mock_backend.delete_collection.call_args_list == [ + call(str(palace_dir), "mempalace_drawers__repair_tmp"), + call(str(palace_dir), "mempalace_drawers"), + call(str(palace_dir), "mempalace_drawers__repair_tmp"), + ] @patch("mempalace.cli.MempalaceConfig") @@ -775,6 +887,7 @@ def test_cmd_repair_aborts_without_confirmation(mock_config_cls, tmp_path, capsy palace_dir.mkdir() (palace_dir / "chroma.sqlite3").write_text("db") mock_config_cls.return_value.palace_path = str(palace_dir) + mock_config_cls.return_value.collection_name = "mempalace_drawers" args = argparse.Namespace(palace=None) mock_col = MagicMock() mock_col.count.return_value = 1 diff --git a/tests/test_config.py b/tests/test_config.py index d7707d9829..90841944ef 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3,7 +3,13 @@ import tempfile import pytest -from mempalace.config import MempalaceConfig, normalize_wing_name, sanitize_kg_value, sanitize_name +from mempalace.config import ( + MempalaceConfig, + normalize_wing_name, + sanitize_kg_value, + sanitize_name, + validate_iso_date, +) def test_default_config(): @@ -212,3 +218,133 @@ def test_kg_value_rejects_null_bytes(): def test_kg_value_rejects_over_length(): with pytest.raises(ValueError): sanitize_kg_value("a" * 129) + + +# --- validate_iso_date --- + + +def test_validate_iso_date_accepts_none(): + assert validate_iso_date(None) is None + + +def test_validate_iso_date_accepts_empty_string(): + assert validate_iso_date("") is None + + +def test_validate_iso_date_accepts_whitespace_only(): + assert validate_iso_date(" ") is None + + +def test_validate_iso_date_accepts_full_date(): + assert validate_iso_date("2025-01-01") == "2025-01-01" + + +def test_validate_iso_date_accepts_year_month(): + assert validate_iso_date("2025-01") == "2025-01" + + +def test_validate_iso_date_accepts_year_only(): + assert validate_iso_date("2025") == "2025" + + +def test_validate_iso_date_strips_whitespace(): + assert validate_iso_date(" 2025-06-15 ") == "2025-06-15" + + +def test_validate_iso_date_rejects_natural_language(): + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("March 2026") + + +def test_validate_iso_date_rejects_partial_date(): + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("2025-1") + + +def test_validate_iso_date_rejects_invalid_month(): + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("2025-13-01") + + +def test_validate_iso_date_rejects_invalid_day(): + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("2025-02-32") + + +def test_validate_iso_date_rejects_garbage(): + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("not-a-date") + + +def test_validate_iso_date_custom_param_name(): + with pytest.raises(ValueError, match="as_of="): + validate_iso_date("Jan 2025", "as_of") + + +def test_validate_iso_date_edge_case_december(): + assert validate_iso_date("2025-12-31") == "2025-12-31" + + +def test_validate_iso_date_edge_case_january(): + assert validate_iso_date("2025-01-01") == "2025-01-01" + + +def test_validate_iso_date_leap_year(): + assert validate_iso_date("2024-02-29") == "2024-02-29" + + +def test_validate_iso_date_non_leap_year_feb_29(): + """2025 is not a leap year — Feb 29 should still pass regex (no calendar check).""" + # The regex validates format, not calendar correctness. + # This is intentional — calendar validation is out of scope. + assert validate_iso_date("2025-02-29") == "2025-02-29" + + +def test_validate_iso_date_rejects_slash_format(): + """YYYY/MM/DD is not ISO-8601.""" + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("2025/01/01") + + +def test_validate_iso_date_rejects_dot_format(): + """DD.MM.YYYY is not ISO-8601.""" + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("01.01.2025") + + +def test_validate_iso_date_rejects_time_component(): + """ISO-8601 datetime with time should be rejected (date only).""" + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("2025-01-01T12:00:00") + + +def test_validate_iso_date_rejects_month_zero(): + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("2025-00-01") + + +def test_validate_iso_date_rejects_day_zero(): + with pytest.raises(ValueError, match="not a valid ISO-8601"): + validate_iso_date("2025-01-00") + + +def test_validate_iso_date_accepts_max_month_day(): + assert validate_iso_date("2025-12-31") == "2025-12-31" + + +def test_validate_iso_date_non_string_input(): + """Non-string input should return None (treated as missing).""" + assert validate_iso_date(123) is None + assert validate_iso_date(3.14) is None + + +def test_validate_iso_date_param_name_in_error(): + """Custom param name appears in the error message.""" + with pytest.raises(ValueError, match="ended="): + validate_iso_date("bad-date", "ended") + + +def test_validate_iso_date_param_name_default(): + """Default param name is 'date'.""" + with pytest.raises(ValueError, match="date="): + validate_iso_date("bad-date") diff --git a/tests/test_convo_miner_unit.py b/tests/test_convo_miner_unit.py index 97236dffb7..b4870eb3f0 100644 --- a/tests/test_convo_miner_unit.py +++ b/tests/test_convo_miner_unit.py @@ -115,6 +115,85 @@ def test_scan_empty_dir(self, tmp_path): files = scan_convos(str(tmp_path)) assert files == [] + def test_scan_skips_subagent_dirs_by_default(self, tmp_path): + # Mimic Claude Code layout: ~/.claude/projects///subagents/agent-*.jsonl + session_dir = tmp_path / "session-abc" + session_dir.mkdir() + (session_dir / "main.jsonl").write_text('{"type":"user"}\n', encoding="utf-8") + subagents_dir = session_dir / "subagents" + subagents_dir.mkdir() + (subagents_dir / "agent-abc.jsonl").write_text('{"type":"user"}\n', encoding="utf-8") + (subagents_dir / "agent-def.jsonl").write_text('{"type":"user"}\n', encoding="utf-8") + + files = scan_convos(str(tmp_path)) + names = [f.name for f in files] + + assert "main.jsonl" in names + assert "agent-abc.jsonl" not in names + assert "agent-def.jsonl" not in names + + def test_scan_includes_subagent_dirs_when_opted_in(self, tmp_path): + session_dir = tmp_path / "session-abc" + session_dir.mkdir() + (session_dir / "main.jsonl").write_text('{"type":"user"}\n', encoding="utf-8") + subagents_dir = session_dir / "subagents" + subagents_dir.mkdir() + (subagents_dir / "agent-abc.jsonl").write_text('{"type":"user"}\n', encoding="utf-8") + + files = scan_convos(str(tmp_path), include_subagents=True) + names = [f.name for f in files] + + assert "main.jsonl" in names + assert "agent-abc.jsonl" in names + + def test_scan_skips_subagent_dirs_at_any_depth(self, tmp_path): + # The "subagents" name match is by directory name, not by depth: verify + # both shallow (top-level) and nested subagents/ get skipped. + (tmp_path / "subagents").mkdir() + (tmp_path / "subagents" / "agent-top.jsonl").write_text("{}", encoding="utf-8") + nested = tmp_path / "session" / "subagents" + nested.mkdir(parents=True) + (nested / "agent-deep.jsonl").write_text("{}", encoding="utf-8") + (tmp_path / "session" / "main.jsonl").write_text("{}", encoding="utf-8") + + files = scan_convos(str(tmp_path)) + names = [f.name for f in files] + + assert "main.jsonl" in names + assert "agent-top.jsonl" not in names + assert "agent-deep.jsonl" not in names + + def test_scan_does_not_skip_suffix_named_dirs(self, tmp_path): + # Exact name match only: 'mysubagents' or 'subagentsbackup' must still + # be mined. Guards against future regression to substring/regex match. + for dir_name in ("mysubagents", "subagentsbackup", "subagent"): + d = tmp_path / dir_name + d.mkdir() + (d / f"{dir_name}.jsonl").write_text("{}", encoding="utf-8") + + files = scan_convos(str(tmp_path)) + names = {f.name for f in files} + + assert "mysubagents.jsonl" in names + assert "subagentsbackup.jsonl" in names + assert "subagent.jsonl" in names + + def test_scan_skips_subagents_case_insensitive(self, tmp_path): + # On Windows + macOS APFS the filesystem is case-preserving; if Claude + # Code or a plugin ever emits 'Subagents' (capitalized), the filter + # must still match. Only one variant per tmp_path because case- + # insensitive filesystems collapse 'Subagents' and 'SUBAGENTS'. + d = tmp_path / "Subagents" + d.mkdir() + (d / "agent.jsonl").write_text("{}", encoding="utf-8") + (tmp_path / "main.jsonl").write_text("{}", encoding="utf-8") + + files = scan_convos(str(tmp_path)) + names = {f.name for f in files} + + assert "main.jsonl" in names + assert "agent.jsonl" not in names + class TestFileChunksLocked: def test_uses_bounded_upsert_batches(self, monkeypatch): diff --git a/tests/test_entity_registry.py b/tests/test_entity_registry.py index c857a071b4..a5f237c983 100644 --- a/tests/test_entity_registry.py +++ b/tests/test_entity_registry.py @@ -2,6 +2,8 @@ from unittest.mock import patch +import pytest + from mempalace.entity_registry import ( COMMON_ENGLISH_WORDS, PERSON_CONTEXT_PATTERNS, @@ -71,6 +73,50 @@ def test_save_creates_file(tmp_path): assert (tmp_path / "entity_registry.json").exists() +def test_save_is_atomic_does_not_leave_tmp(tmp_path): + # Atomic write must not leave the .tmp sidecar file after a successful save. + registry = EntityRegistry.load(config_dir=tmp_path) + registry.save() + leftover = list(tmp_path.glob("entity_registry.json.tmp*")) + assert leftover == [], f"atomic write leaked tmp file(s): {leftover}" + + +def test_save_preserves_previous_on_serialization_failure(tmp_path, monkeypatch): + # If serialization fails mid-write, the previous registry must remain + # intact — this is the whole point of atomic write vs truncating in place. + registry = EntityRegistry.load(config_dir=tmp_path) + registry.seed( + mode="personal", + people=[{"name": "Alice", "relationship": "friend", "context": "personal"}], + projects=[], + ) + registry.save() + target = tmp_path / "entity_registry.json" + original = target.read_text(encoding="utf-8") + + # Force os.replace to raise — simulates filesystem full / permission flip + # AFTER the temp file is written but BEFORE the rename completes. + import os as _os + + real_replace = _os.replace + + def boom(src, dst): + raise OSError("simulated rename failure") + + monkeypatch.setattr(_os, "replace", boom) + with pytest.raises(OSError): + registry.seed( + mode="personal", + people=[{"name": "Bob", "relationship": "friend", "context": "personal"}], + projects=[], + ) + registry.save() + + # Restore os.replace before reading so the assertion can rely on it. + monkeypatch.setattr(_os, "replace", real_replace) + assert target.read_text(encoding="utf-8") == original + + # ── seed ──────────────────────────────────────────────────────────────── diff --git a/tests/test_hnsw_capacity.py b/tests/test_hnsw_capacity.py index 512fc9c677..912def8bf7 100644 --- a/tests/test_hnsw_capacity.py +++ b/tests/test_hnsw_capacity.py @@ -238,14 +238,39 @@ def test_capacity_status_tolerates_flush_lag(tmp_path): assert info["status"] == "ok" -def test_capacity_status_flags_unflushed_with_large_sqlite(tmp_path): - """No pickle + many sqlite rows is its own divergence signal.""" +def test_capacity_status_does_not_flag_unflushed_with_large_sqlite(tmp_path): + """No pickle + many sqlite rows is inconclusive, not divergence.""" seg = "seg-noflush" _seed_chroma_db(str(tmp_path), sqlite_count=10_000, segment_id=seg) info = hnsw_capacity_status(str(tmp_path), COLLECTION) - assert info["diverged"] is True + assert info["diverged"] is False + assert info["status"] == "unknown" + assert info["divergence"] is None assert info["hnsw_count"] is None - assert "never flushed" in info["message"] + assert "capacity unavailable" in info["message"] + assert "leaving vector search enabled" in info["message"] + + +def test_mcp_probe_does_not_disable_vectors_for_unflushed_metadata(tmp_path, monkeypatch): + """The MCP preflight must not route all searches to BM25 on this signal.""" + from mempalace import mcp_server + + seg = "seg-mcp-noflush" + _seed_chroma_db(str(tmp_path), sqlite_count=10_000, segment_id=seg) + + class _Cfg: + palace_path = str(tmp_path) + + monkeypatch.setattr(mcp_server, "_config", _Cfg()) + monkeypatch.setattr(mcp_server, "_vector_disabled", True) + monkeypatch.setattr(mcp_server, "_vector_disabled_reason", "old divergence") + + mcp_server._refresh_vector_disabled_flag() + + assert mcp_server._vector_disabled is False + assert mcp_server._vector_disabled_reason == "" + assert mcp_server._vector_capacity_status["status"] == "unknown" + assert "leaving vector search enabled" in mcp_server._vector_capacity_status["message"] def test_capacity_status_quiet_for_empty_palace(tmp_path): @@ -372,6 +397,17 @@ def _seed_drawers(palace: str, segment_id: str, drawers: list[tuple[str, dict, s conn.close() +def _set_drawer_created_at(palace: str, timestamps: dict[int, str]) -> None: + db_path = os.path.join(palace, "chroma.sqlite3") + conn = sqlite3.connect(db_path) + try: + for emb_id, created_at in timestamps.items(): + conn.execute("UPDATE embeddings SET created_at = ? WHERE id = ?", (created_at, emb_id)) + conn.commit() + finally: + conn.close() + + @pytest.fixture def palace_with_drawers(tmp_path): seg = "seg-bm25" @@ -417,6 +453,122 @@ def test_bm25_fallback_filters_by_wing(palace_with_drawers): assert all(r["wing"] == "design" for r in out["results"]) +def test_bm25_fallback_applies_wing_before_fts_candidate_limit(tmp_path): + seg = "seg-bm25-fts-limit" + _seed_chroma_db(str(tmp_path), sqlite_count=0, segment_id=seg) + _seed_drawers( + str(tmp_path), + seg, + [ + ( + "shared token outside target wing", + {"wing": "ops", "room": "incidents", "source_file": "/x/ops.md"}, + "d-1", + ), + ( + "shared token inside target wing", + {"wing": "project", "room": "diary", "source_file": "/x/project.md"}, + "d-2", + ), + ], + ) + + out = _bm25_only_via_sqlite("shared token", str(tmp_path), wing="project", max_candidates=1) + + assert out["total_before_filter"] == 1 + assert len(out["results"]) == 1 + assert out["results"][0]["wing"] == "project" + + +def test_bm25_fallback_applies_room_before_fts_candidate_limit(tmp_path): + seg = "seg-bm25-room-limit" + _seed_chroma_db(str(tmp_path), sqlite_count=0, segment_id=seg) + _seed_drawers( + str(tmp_path), + seg, + [ + ( + "shared token wrong room", + {"wing": "project", "room": "scratch", "source_file": "/x/scratch.md"}, + "d-1", + ), + ( + "shared token right room", + {"wing": "project", "room": "diary", "source_file": "/x/diary.md"}, + "d-2", + ), + ], + ) + + out = _bm25_only_via_sqlite( + "shared token", + str(tmp_path), + wing="project", + room="diary", + max_candidates=1, + ) + + assert out["total_before_filter"] == 1 + assert len(out["results"]) == 1 + assert out["results"][0]["wing"] == "project" + assert out["results"][0]["room"] == "diary" + + +def test_bm25_fallback_applies_wing_before_recency_candidate_limit(tmp_path): + seg = "seg-bm25-recency-limit" + _seed_chroma_db(str(tmp_path), sqlite_count=0, segment_id=seg) + _seed_drawers( + str(tmp_path), + seg, + [ + ( + "target drawer for short query", + {"wing": "project", "room": "diary", "source_file": "/x/project.md"}, + "d-1", + ), + ( + "newer drawer outside target wing", + {"wing": "ops", "room": "incidents", "source_file": "/x/ops.md"}, + "d-2", + ), + ], + ) + _set_drawer_created_at( + str(tmp_path), + { + 1: "2026-01-01 00:00:00", + 2: "2026-02-01 00:00:00", + }, + ) + + out = _bm25_only_via_sqlite("a", str(tmp_path), wing="project", max_candidates=1) + + assert out["total_before_filter"] == 1 + assert len(out["results"]) == 1 + assert out["results"][0]["wing"] == "project" + + +def test_bm25_fallback_returns_empty_when_filtered_wing_has_no_candidates(tmp_path): + seg = "seg-bm25-empty-filter" + _seed_chroma_db(str(tmp_path), sqlite_count=0, segment_id=seg) + _seed_drawers( + str(tmp_path), + seg, + [ + ( + "shared token outside target wing", + {"wing": "ops", "room": "incidents", "source_file": "/x/ops.md"}, + "d-1", + ), + ], + ) + + out = _bm25_only_via_sqlite("shared token", str(tmp_path), wing="project", max_candidates=1) + + assert out["total_before_filter"] == 0 + assert out["results"] == [] + + def test_bm25_fallback_no_palace(tmp_path): out = _bm25_only_via_sqlite("anything", str(tmp_path)) assert "error" in out diff --git a/tests/test_hooks_cli.py b/tests/test_hooks_cli.py index 1ceb530ffb..a72ef5db48 100644 --- a/tests/test_hooks_cli.py +++ b/tests/test_hooks_cli.py @@ -13,6 +13,7 @@ _count_human_messages, _extract_recent_messages, _get_mine_targets, + _has_meaningful_updates, _log, _maybe_auto_ingest, _mempalace_python, @@ -161,6 +162,18 @@ def test_extract_recent_messages_missing_file(): assert _extract_recent_messages("/nonexistent.jsonl") == [] +def test_meaningful_update_heuristic_skips_chatter(): + assert _has_meaningful_updates(["continue", "ok", "thanks", "yes"]) is False + + +def test_meaningful_update_heuristic_detects_real_work(): + assert _has_meaningful_updates([ + "Audit aws infra and confirm the ananas-hub RDS deployment status.", + "Fix the marketing route and replace the placeholder API with a real dataset query.", + "Document the architecture decision in /home/zapostolski/projects/ananas-hub/docs/architecture.md.", + ]) is True + + # --- hook_stop --- @@ -221,10 +234,15 @@ def test_stop_hook_passthrough_below_interval(tmp_path): def test_stop_hook_saves_silently_at_interval(tmp_path): transcript = tmp_path / "t.jsonl" - _write_transcript( - transcript, - [{"message": {"role": "user", "content": f"msg {i}"}} for i in range(SAVE_INTERVAL)], - ) + _write_transcript(transcript, [ + { + "message": { + "role": "user", + "content": f"Fix marketing data contract {i} and audit the Power BI dataset wiring.", + } + } + for i in range(SAVE_INTERVAL) + ]) save_result = {"count": 15, "themes": ["hooks", "notifications"]} with patch("mempalace.hooks_cli._save_diary_direct", return_value=save_result) as mock_save: result = _capture_hook_output( @@ -246,7 +264,15 @@ def test_stop_hook_derives_wing_from_transcript_path(tmp_path): transcript = project_dir / "session.jsonl" _write_transcript( transcript, - [{"message": {"role": "user", "content": f"msg {i}"}} for i in range(SAVE_INTERVAL)], + [ + { + "message": { + "role": "user", + "content": f"Implement aws audit step {i} and update /home/zapostolski/projects/ananas-hub/README.md.", + } + } + for i in range(SAVE_INTERVAL) + ], ) save_result = {"count": 15, "themes": []} with patch("mempalace.hooks_cli._save_diary_direct", return_value=save_result) as mock_save: @@ -258,12 +284,31 @@ def test_stop_hook_derives_wing_from_transcript_path(tmp_path): mock_save.assert_called_once_with(str(transcript), "test", wing="wing_myproject", toast=False) -def test_stop_hook_tracks_save_point(tmp_path): +def test_stop_hook_skips_low_signal_at_interval(tmp_path): transcript = tmp_path / "t.jsonl" - _write_transcript( - transcript, - [{"message": {"role": "user", "content": f"msg {i}"}} for i in range(SAVE_INTERVAL)], + _write_transcript(transcript, [ + {"message": {"role": "user", "content": "continue"}} + for _ in range(SAVE_INTERVAL) + ]) + result = _capture_hook_output( + hook_stop, + {"session_id": "test", "stop_hook_active": False, "transcript_path": str(transcript)}, + state_dir=tmp_path, ) + assert result == {} + + +def test_stop_hook_tracks_save_point(tmp_path): + transcript = tmp_path / "t.jsonl" + _write_transcript(transcript, [ + { + "message": { + "role": "user", + "content": f"Implement aws audit step {i} and update /home/zapostolski/projects/ananas-hub/README.md.", + } + } + for i in range(SAVE_INTERVAL) + ]) data = {"session_id": "test", "stop_hook_active": False, "transcript_path": str(transcript)} # First call saves silently with systemMessage notification @@ -282,7 +327,32 @@ def test_stop_hook_tracks_save_point(tmp_path): # --- hook_session_start --- -def test_session_start_passes_through(tmp_path): +def test_session_start_injects_context(tmp_path, monkeypatch): + """SessionStart returns Claude-compatible additionalContext from wake-up + protocol nudge.""" + import mempalace.hooks_cli as hc + monkeypatch.setattr(hc, "_build_session_start_context", lambda cwd, palace: "WAKE\n\nNUDGE") + + result = _capture_hook_output( + hook_session_start, + {"session_id": "test", "cwd": "/tmp"}, + state_dir=tmp_path, + ) + out = result["hookSpecificOutput"] + assert out["hookEventName"] == "SessionStart" + assert "WAKE" in out["additionalContext"] + assert "NUDGE" in out["additionalContext"] + + +def test_session_start_skips_when_palace_unresolvable(tmp_path, monkeypatch): + """If palace path resolution fails, SessionStart is a no-op (avoid breaking the harness).""" + import mempalace.hooks_cli as hc + import mempalace.config as config_mod + + def _boom(*a, **kw): + raise RuntimeError("no palace configured") + + monkeypatch.setattr(config_mod, "MempalaceConfig", _boom) + result = _capture_hook_output( hook_session_start, {"session_id": "test"}, @@ -291,6 +361,145 @@ def test_session_start_passes_through(tmp_path): assert result == {} +# --- _mirror_local_memory --- + + +def _make_md(path: Path, body: str = "# decision\n\nUse Postgres for ledger.\n"): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(body, encoding="utf-8") + + +def test_mirror_skips_index_files(tmp_path, monkeypatch): + """MEMORY.md / CLAUDE.md / GEMINI.md are skipped — they're indexes, not content.""" + import mempalace.hooks_cli as hc + root = tmp_path / "claude" / "projects" + _make_md(root / "-home-zap-projects-foo" / "MEMORY.md", "- pointer\n") + _make_md(root / "-home-zap-projects-foo" / "CLAUDE.md", "context\n") + _make_md(root / "-home-zap-projects-foo" / "memory" / "real.md") + + seen_calls = [] + + def fake_add(wing, room, content, source_file=None, added_by="mcp"): + seen_calls.append((wing, room)) + return {"success": True, "drawer_id": f"d_{wing}_{room}"} + + monkeypatch.setattr(hc, "STATE_DIR", tmp_path / "state") + monkeypatch.setattr(hc, "MIRROR_STATE_FILE", tmp_path / "state" / "mirror.json") + monkeypatch.setattr(hc, "DEFAULT_MIRROR_ROOTS", (root,)) + monkeypatch.setattr("mempalace.mcp_server.tool_add_drawer", fake_add) + + counts = hc._mirror_local_memory() + assert counts["added"] == 1 + assert seen_calls == [("foo", "real")] + + +def test_mirror_idempotent_on_unchanged(tmp_path, monkeypatch): + """Second run with no changes adds nothing.""" + import mempalace.hooks_cli as hc + root = tmp_path / "claude" / "projects" + _make_md(root / "-home-zap-projects-bar" / "memory" / "decision.md") + + monkeypatch.setattr(hc, "STATE_DIR", tmp_path / "state") + monkeypatch.setattr(hc, "MIRROR_STATE_FILE", tmp_path / "state" / "mirror.json") + monkeypatch.setattr(hc, "DEFAULT_MIRROR_ROOTS", (root,)) + monkeypatch.setattr( + "mempalace.mcp_server.tool_add_drawer", + lambda **kw: {"success": True, "drawer_id": "d1"}, + ) + + first = hc._mirror_local_memory() + second = hc._mirror_local_memory() + assert first["added"] == 1 + assert second["added"] == 0 + assert second["skipped_unchanged"] == 1 + + +def test_mirror_re_adds_when_mtime_changes(tmp_path, monkeypatch): + """Editing a file bumps mtime → re-mirrored.""" + import mempalace.hooks_cli as hc + root = tmp_path / "claude" / "projects" + md = root / "-home-zap-projects-baz" / "memory" / "plan.md" + _make_md(md) + + add_calls = [] + monkeypatch.setattr(hc, "STATE_DIR", tmp_path / "state") + monkeypatch.setattr(hc, "MIRROR_STATE_FILE", tmp_path / "state" / "mirror.json") + monkeypatch.setattr(hc, "DEFAULT_MIRROR_ROOTS", (root,)) + monkeypatch.setattr( + "mempalace.mcp_server.tool_add_drawer", + lambda **kw: add_calls.append(kw) or {"success": True, "drawer_id": "d"}, + ) + + hc._mirror_local_memory() + # bump mtime + os.utime(md, (md.stat().st_atime, md.stat().st_mtime + 10)) + hc._mirror_local_memory() + assert len(add_calls) == 2 + + +def test_mirror_handles_dup_response(tmp_path, monkeypatch): + """When mempalace says duplicate, we still record state to avoid retry.""" + import mempalace.hooks_cli as hc + root = tmp_path / "claude" / "projects" + _make_md(root / "-home-zap-projects-qux" / "memory" / "x.md") + + monkeypatch.setattr(hc, "STATE_DIR", tmp_path / "state") + monkeypatch.setattr(hc, "MIRROR_STATE_FILE", tmp_path / "state" / "mirror.json") + monkeypatch.setattr(hc, "DEFAULT_MIRROR_ROOTS", (root,)) + monkeypatch.setattr( + "mempalace.mcp_server.tool_add_drawer", + lambda **kw: {"success": False, "reason": "duplicate", "matches": []}, + ) + + first = hc._mirror_local_memory() + second = hc._mirror_local_memory() + assert first["skipped_dup"] == 1 + assert second["skipped_unchanged"] == 1 + + +def test_derive_wing_room_claude_slug(): + import mempalace.hooks_cli as hc + root = Path("/home/u/.claude/projects") + p = root / "-home-u-projects-ai-marketing" / "memory" / "decisions.md" + assert hc._derive_wing_room(p, root) == ("ai-marketing", "decisions") + + +def test_derive_wing_room_gemini_layout(): + import mempalace.hooks_cli as hc + root = Path("/home/u/.gemini/memory") + p = root / "ananas-crm" / "GEMINI.md" # GEMINI.md filtered earlier; testing derivation only + assert hc._derive_wing_room(p, root) == ("ananas-crm", "GEMINI") + + +def test_mirror_disabled_via_env(tmp_path, monkeypatch): + """MEMPAL_MIRROR_DISABLED=1 in hook_stop short-circuits the mirror.""" + import mempalace.hooks_cli as hc + root = tmp_path / "claude" / "projects" + _make_md(root / "-home-zap-projects-skip" / "memory" / "x.md") + + monkeypatch.setenv("MEMPAL_MIRROR_DISABLED", "1") + monkeypatch.setattr(hc, "STATE_DIR", tmp_path / "state") + monkeypatch.setattr(hc, "MIRROR_STATE_FILE", tmp_path / "state" / "mirror.json") + monkeypatch.setattr(hc, "DEFAULT_MIRROR_ROOTS", (root,)) + + called = {"n": 0} + + def boom(*a, **kw): + called["n"] += 1 + raise AssertionError("mirror should not have run") + + monkeypatch.setattr(hc, "_mirror_local_memory", boom) + + transcript = tmp_path / "t.jsonl" + _write_transcript(transcript, [{"message": {"role": "user", "content": "ok"}}]) + _capture_hook_output( + hook_stop, + {"session_id": "test", "stop_hook_active": False, "transcript_path": str(transcript)}, + state_dir=tmp_path, + ) + assert called["n"] == 0 + + # --- hook_precompact --- @@ -504,7 +713,7 @@ def test_maybe_auto_ingest_ignores_transcript_arg_path(tmp_path): Transcript convos are handled by _ingest_transcript (called separately in hook handlers). _maybe_auto_ingest only handles MEMPAL_DIR — even when invoked in a context where a transcript is also being processed, - no second spawn for the transcript dir should appear here. + no second_spawn for the transcript dir should appear here. """ convo_dir = tmp_path / "convos" convo_dir.mkdir() diff --git a/tests/test_knowledge_graph.py b/tests/test_knowledge_graph.py index d7d9838fa4..6eeb8d3f5f 100644 --- a/tests/test_knowledge_graph.py +++ b/tests/test_knowledge_graph.py @@ -5,6 +5,8 @@ timeline, stats, and edge cases (duplicate triples, ID collisions). """ +import pytest + class TestEntityOperations: def test_add_entity(self, kg): @@ -45,6 +47,38 @@ def test_invalidated_triple_allows_re_add(self, kg): tid2 = kg.add_triple("Alice", "works_at", "Acme") assert tid1 != tid2 # new triple since old one was closed + def test_add_triple_rejects_inverted_interval(self, kg): + # valid_to before valid_from would never satisfy + # `valid_from <= as_of AND valid_to >= as_of` — silently invisible + # to every query. Reject at write time instead. + with pytest.raises(ValueError, match="before valid_from"): + kg.add_triple( + "Alice", + "worked_at", + "Acme", + valid_from="2026-03-01", + valid_to="2026-02-01", + ) + + def test_add_triple_accepts_equal_dates(self, kg): + # Same-day intervals are valid (point-in-time facts). + tid = kg.add_triple( + "Alice", + "joined", + "Acme", + valid_from="2026-03-15", + valid_to="2026-03-15", + ) + assert tid.startswith("t_alice_joined_acme_") + + def test_add_triple_allows_only_one_bound(self, kg): + # The guard only fires when BOTH bounds are set. + tid1 = kg.add_triple("Alice", "knows", "Bob", valid_from="2026-01-01") + assert tid1.startswith("t_alice_knows_bob_") + kg.invalidate("Alice", "knows", "Bob", ended="2026-02-01") + tid2 = kg.add_triple("Alice", "knew", "Bob", valid_to="2026-03-01") + assert tid2.startswith("t_alice_knew_bob_") + class TestQueries: def test_query_outgoing(self, seeded_kg): diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index ee0fa9c201..116f2191b4 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -457,6 +457,26 @@ def test_add_drawer_duplicate_detection(self, monkeypatch, config, palace_path, assert result2["success"] is True assert result2["reason"] == "already_exists" + def test_add_drawer_fails_when_readback_misses(self, monkeypatch, config, kg): + _patch_mcp_server(monkeypatch, config, kg) + from mempalace import mcp_server + + class _FakeGetResult: + ids = [] + + class _FakeCol: + def get(self, **kwargs): + return _FakeGetResult() + + def upsert(self, **kwargs): + return None + + monkeypatch.setattr(mcp_server, "_get_collection", lambda create=False: _FakeCol()) + + result = mcp_server.tool_add_drawer("w", "r", "content") + assert result["success"] is False + assert "not readable" in result["error"] + def test_add_drawer_shared_header_no_collision(self, monkeypatch, config, palace_path, kg): """Documents sharing a >100-char header must get distinct IDs (full-content hash).""" _patch_mcp_server(monkeypatch, config, kg) @@ -476,9 +496,9 @@ def test_add_drawer_shared_header_no_collision(self, monkeypatch, config, palace assert result1["success"] is True assert result2["success"] is True - assert ( - result1["drawer_id"] != result2["drawer_id"] - ), "Documents with shared header but different content must have distinct drawer IDs" + assert result1["drawer_id"] != result2["drawer_id"], ( + "Documents with shared header but different content must have distinct drawer IDs" + ) def test_delete_drawer(self, monkeypatch, config, palace_path, seeded_collection, kg): _patch_mcp_server(monkeypatch, config, kg) @@ -669,6 +689,90 @@ def test_kg_invalidate(self, monkeypatch, config, palace_path, seeded_kg): ended="2026-03-01", ) assert result["success"] is True + # Regression #1314: response must echo the actual ended date, + # not silently drop it and return the literal string "today". + assert result["ended"] == "2026-03-01" + + def test_kg_add_forwards_valid_to(self, monkeypatch, config, palace_path, kg): + """Regression #1314 case 1: valid_to must round-trip through kg_add.""" + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_kg_add + + result = tool_kg_add( + subject="_test_temporal", + predicate="had_value", + object="probe", + valid_from="2026-01-01", + valid_to="2026-04-28", + ) + assert result["success"] is True + + facts = kg.query_entity("_test_temporal") + assert len(facts) == 1 + assert facts[0]["valid_from"] == "2026-01-01" + assert facts[0]["valid_to"] == "2026-04-28" + # An already-ended fact must not be reported as still current. + assert facts[0]["current"] is False + + def test_kg_add_forwards_source_provenance(self, monkeypatch, config, palace_path, kg): + """Regression #1314 case 3: source_file / source_drawer_id reach storage.""" + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_kg_add + + result = tool_kg_add( + subject="operating-verb", + predicate="candidate", + object="husbandry", + valid_from="2026-04-28", + source_closet="closet-42", + source_file="docs/decisions.md", + source_drawer_id="drawer_abc123", + ) + assert result["success"] is True + + triple_id = result["triple_id"] + # Read raw row to verify all provenance columns persisted. + with kg._lock: + row = ( + kg._conn() + .execute( + "SELECT source_closet, source_file, source_drawer_id FROM triples WHERE id = ?", + (triple_id,), + ) + .fetchone() + ) + assert row is not None + assert row["source_closet"] == "closet-42" + assert row["source_file"] == "docs/decisions.md" + assert row["source_drawer_id"] == "drawer_abc123" + + def test_kg_invalidate_returns_actual_ended_date( + self, monkeypatch, config, palace_path, seeded_kg + ): + """Regression #1314 case 2: response reports the resolved date, not 'today'.""" + from datetime import date as _date + + _patch_mcp_server(monkeypatch, config, seeded_kg) + from mempalace.mcp_server import tool_kg_invalidate + + # Caller-supplied date round-trips into the response. + explicit = tool_kg_invalidate( + subject="Max", + predicate="does", + object="swimming", + ended="2026-04-28", + ) + assert explicit["ended"] == "2026-04-28" + + # Caller-omitted date resolves to today's ISO date — never the + # literal string "today" the buggy implementation used to return. + implicit = tool_kg_invalidate( + subject="Max", + predicate="loves", + object="Chess", + ) + assert implicit["ended"] != "today" + assert implicit["ended"] == _date.today().isoformat() def test_kg_timeline(self, monkeypatch, config, palace_path, seeded_kg): _patch_mcp_server(monkeypatch, config, seeded_kg) @@ -701,7 +805,8 @@ def test_diary_write_and_read(self, monkeypatch, config, palace_path, kg): topic="architecture", ) assert w["success"] is True - assert w["agent"] == "TestAgent" + # agent_name is normalized to lowercase on write (#1243). + assert w["agent"] == "testagent" r = tool_diary_read(agent_name="TestAgent") assert r["total"] == 1 @@ -793,6 +898,50 @@ def test_diary_read_empty_wing_spans_all_wings(self, monkeypatch, config, palace assert r_scoped["total"] == 1 assert r_scoped["entries"][0]["content"] == "project-wing entry" + def test_diary_read_case_insensitive_agent(self, monkeypatch, config, palace_path, kg): + """Regression for #1243: diary_read must be case-insensitive over + agent_name. Writing as "Claude" and reading as "claude" (or vice + versa) must surface the same entries — sanitize_name preserved + case, which silently dropped reads when the agent name's casing + differed from the write.""" + _patch_mcp_server(monkeypatch, config, kg) + _client, _col = _get_collection(palace_path, create=True) + del _client + from mempalace.mcp_server import tool_diary_read, tool_diary_write + + # Write as "Claude" → read as "claude" should match. + w1 = tool_diary_write( + agent_name="Claude", + entry="entry written as Claude", + topic="general", + ) + assert w1["success"] + + r1 = tool_diary_read(agent_name="claude") + assert "entries" in r1, r1 + contents1 = {e["content"] for e in r1["entries"]} + assert "entry written as Claude" in contents1 + + # Write as "CLAUDE" → read as "Claude" should also match the + # same agent. After normalization both writes target the same + # lowercase agent identity, so both entries are returned. + w2 = tool_diary_write( + agent_name="CLAUDE", + entry="entry written as CLAUDE", + topic="general", + ) + assert w2["success"] + + r2 = tool_diary_read(agent_name="Claude") + contents2 = {e["content"] for e in r2["entries"]} + assert "entry written as Claude" in contents2 + assert "entry written as CLAUDE" in contents2 + + # The stored agent metadata is the lowercase form, and the + # default wing is derived from that lowercase form too. + assert w1["agent"] == "claude" + assert w2["agent"] == "claude" + # ── Cache Invalidation (inode/mtime) ────────────────────────────────── @@ -894,6 +1043,25 @@ def test_reconnect_reports_success(self, monkeypatch, config, palace_path, kg): assert "Reconnected" in result["message"] assert isinstance(result["drawers"], int) + def test_reconnect_closes_shared_backend(self, monkeypatch, config, kg): + _patch_mcp_server(monkeypatch, config, kg) + from unittest.mock import MagicMock + + from mempalace import mcp_server, palace + + close_palace = MagicMock() + monkeypatch.setattr(palace._DEFAULT_BACKEND, "close_palace", close_palace) + + class _FakeCol: + def count(self): + return 7 + + monkeypatch.setattr(mcp_server, "_get_collection", lambda create=False: _FakeCol()) + + result = mcp_server.tool_reconnect() + assert result["success"] is True + close_palace.assert_called_once_with(config.palace_path) + def test_get_collection_create_true_avoids_get_or_create_on_reopen( self, monkeypatch, config, palace_path, kg ): @@ -938,3 +1106,59 @@ def _spy(self, *args, **kwargs): col2 = mcp_server._get_collection(create=True) assert col2 is not None assert calls == [], f"get_or_create_collection was called: {calls}" + + def test_get_collection_passes_embedding_function(self, monkeypatch, config, palace_path, kg): + """Regression for #1299. + + ``mcp_server._get_collection`` must pass ``embedding_function=`` into + both ``client.get_collection`` and ``client.create_collection``, + mirroring ``ChromaBackend.get_collection``. Without it, ChromaDB 1.x + falls back to its built-in ``DefaultEmbeddingFunction`` (whose lazy + ONNX provider selection has SIGSEGV'd on python 3.14 + Apple Silicon), + and writers/readers can disagree with the miner about which EF is + bound to the collection. The miner / Stop hook ingest path routes + through ``ChromaBackend.get_collection`` which does this correctly; + the MCP server must match. + """ + _patch_mcp_server(monkeypatch, config, kg) + from mempalace import mcp_server + + client = mcp_server._get_client() + client_cls = type(client) + captured: dict[str, list[dict]] = {"get": [], "create": []} + real_get = client_cls.get_collection + real_create = client_cls.create_collection + + def _spy_get(self, name, **kwargs): + captured["get"].append(dict(kwargs)) + return real_get(self, name, **kwargs) + + def _spy_create(self, name, **kwargs): + captured["create"].append(dict(kwargs)) + return real_create(self, name, **kwargs) + + monkeypatch.setattr(client_cls, "get_collection", _spy_get) + monkeypatch.setattr(client_cls, "create_collection", _spy_create) + mcp_server._collection_cache = None + + col = mcp_server._get_collection(create=True) + assert col is not None + + all_calls = captured["get"] + captured["create"] + assert all_calls, "expected get_collection or create_collection to be called" + for kwargs in all_calls: + assert "embedding_function" in kwargs, ( + f"missing embedding_function= in chromadb call: {kwargs}" + ) + assert kwargs["embedding_function"] is not None + + # Same expectation on the create=False (cache-miss) reopen path. + mcp_server._collection_cache = None + captured["get"].clear() + captured["create"].clear() + col2 = mcp_server._get_collection() + assert col2 is not None + assert captured["get"], "expected get_collection on cache-miss reopen" + for kwargs in captured["get"]: + assert "embedding_function" in kwargs + assert kwargs["embedding_function"] is not None diff --git a/tests/test_normalize.py b/tests/test_normalize.py index 2b0f180710..540b332ef1 100644 --- a/tests/test_normalize.py +++ b/tests/test_normalize.py @@ -1316,6 +1316,27 @@ def test_plus_n_lines_marker_inline(self): text = "> User:\n> The log showed … +50 lines of stack trace, useful." assert strip_noise(text) == text.strip() + def test_user_documents_ansi_escape_by_name(self): + # User prose that mentions an ANSI sequence by name (e.g. in docs) + # contains no ESC byte, so the ANSI sweep must not touch it. + text = ( + "> User:\n" + "> The bold SGR code is `[1m` and reset is `[22m`. " + "Documenting the format here." + ) + assert strip_noise(text) == text.strip() + + def test_user_mentions_local_command_inline(self): + # Inline mention of a Claude Code envelope tag inside user prose + # (e.g. when documenting the harness) must not be stripped. + text = ( + "> User:\n" + "> Claude Code wraps slash commands in " + "..." + " — that is the tag name." + ) + assert strip_noise(text) == text.strip() + def test_dangling_open_tag_does_not_span_messages(self): # THE span-eating bug: a stray unclosed in one # message must NOT merge with a closing tag in another message and @@ -1391,15 +1412,92 @@ def test_strips_each_known_noise_tag(self): "system-reminder", "command-message", "command-name", + "command-args", "task-notification", "user-prompt-submit-hook", "hook_output", + "local-command-caveat", + "local-command-stdout", ): text = f"> User:\n<{tag}>junk\n> Real." out = strip_noise(text) assert tag not in out, f"{tag} leaked into output" assert "Real." in out + def test_strips_full_claude_code_slash_command_envelope(self): + # Every slash-command invocation in Claude Code emits this 5-tag + # envelope. Before the fix, only command-name and command-message + # were stripped; the other three survived into stored drawers. + text = ( + "Caveat: do not respond.\n" + "/effort\n" + "effort\n" + "max\n" + "Set effort level to max\n" + "Real follow-up content." + ) + out = strip_noise(text) + for tag in ( + "local-command-caveat", + "command-name", + "command-message", + "command-args", + "local-command-stdout", + ): + assert tag not in out, f"{tag} leaked" + assert "Real follow-up content." in out + + def test_strips_empty_command_args_pair(self): + # Claude Code emits with empty body + # when a slash command takes no arguments. The pattern still matches. + text = "> clear\n" + out = strip_noise(text) + assert "command-args" not in out + assert "command-message" not in out + + def test_strips_ansi_color_codes(self): + # SGR color codes from Bash tool output. Surrounding text preserved. + text = "before \x1b[1mbold\x1b[22m after" + out = strip_noise(text) + assert "\x1b" not in out + assert "before" in out + assert "bold" in out + assert "after" in out + + def test_strips_ansi_truecolor_codes(self): + # 24-bit truecolor sequences from the /context renderer. + text = "\x1b[38;2;153;153;153m├\x1b[39m mempalace_add_drawer" + out = strip_noise(text) + assert "\x1b" not in out + assert "├ mempalace_add_drawer" in out + + def test_strips_ansi_cursor_moves(self): + text = "before\x1b[2K\x1b[Hafter" + out = strip_noise(text) + assert "\x1b" not in out + assert "beforeafter" in out + + def test_strips_ansi_osc_terminal_title(self): + # OSC 0 terminator: BEL. + text = "before\x1b]0;window title\x07after" + out = strip_noise(text) + assert "\x1b" not in out + assert "beforeafter" in out + + def test_strips_ansi_osc_hyperlink(self): + # OSC 8 hyperlinks terminated by ST (ESC \). + text = "see \x1b]8;;https://example.com\x1b\\here\x1b]8;;\x1b\\." + out = strip_noise(text) + assert "\x1b" not in out + assert "see here." in out + + def test_strips_ansi_inside_noise_tag_with_tag(self): + # ANSI inside a tag exits with the tag — the dedicated ANSI sweep + # only has to cover standalone ANSI in tool output, not double-strip. + text = "Set mode to \x1b[1mdefault\x1b[22m" + out = strip_noise(text) + assert out == "" + def test_collapses_excessive_blank_lines(self): text = "line one\n\n\n\n\n\nline two" out = strip_noise(text) diff --git a/tests/test_palace_locks.py b/tests/test_palace_locks.py index 601c8941af..39aa50c597 100644 --- a/tests/test_palace_locks.py +++ b/tests/test_palace_locks.py @@ -135,19 +135,77 @@ def test_different_palaces_dont_conflict(tmp_path, monkeypatch): def test_palace_path_is_normalized(tmp_path, monkeypatch): - """Relative and absolute forms of the same path must use the same lock.""" + """Relative and absolute forms of the same path must use the same lock. + + Cross-process variant: a child holds the absolute form, a relative form + in the parent must hash to the same lock key and raise + ``MineAlreadyRunning``. (The same-thread case is now a re-entrant + pass-through by design — see ``test_reentrant_same_thread_passes_through`` + — so we exercise the normalization invariant across a process boundary + where re-entrance does not apply.) + """ monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.chdir(tmp_path) os.makedirs(tmp_path / "palace", exist_ok=True) absolute = str(tmp_path / "palace") - relative = "palace" + ready = str(tmp_path / "ready") + release = str(tmp_path / "release") - # Hold the lock with the absolute form; attempting to re-acquire with - # the relative form (which resolves to the same absolute path) must fail. - with mine_palace_lock(absolute): + ctx = _get_mp_context() + holder = ctx.Process(target=_hold_lock, args=(absolute, ready, release)) + holder.start() + try: + for _ in range(500): + if os.path.exists(ready): + break + time.sleep(0.01) + assert os.path.exists(ready), "holder failed to acquire lock in time" + + # Parent holds CWD = tmp_path so "palace" is the same on-disk dir as + # the absolute form. The lock key is sha256(realpath+normcase) so the + # two forms must collide. with pytest.raises(MineAlreadyRunning): - with mine_palace_lock(relative): + with mine_palace_lock("palace"): pytest.fail("normalized path collision should have raised") + finally: + open(release, "w").close() + holder.join(timeout=5) + + +def test_reentrant_same_thread_passes_through(tmp_path, monkeypatch): + """Same thread re-acquiring the same palace lock must not deadlock or raise. + + This is the invariant that makes ``ChromaCollection`` write methods (which + take ``mine_palace_lock`` for MCP/direct-writer protection) compose with + ``miner.mine()`` (which already holds the lock for the entire mine + pipeline). Without the per-thread re-entrant guard the inner acquire + would self-deadlock on the outer flock. + """ + monkeypatch.setenv("HOME", str(tmp_path)) + palace = str(tmp_path / "palace") + with mine_palace_lock(palace): + # Re-enter from the same thread — must yield without raising or hanging. + with mine_palace_lock(palace): + pass + # After the inner exits, the outer is still held: confirm via a + # subprocess that tries to acquire and reports back. + ctx = _get_mp_context() + result_q = ctx.Queue() + child = ctx.Process(target=_try_acquire_expect_busy, args=(palace, result_q)) + child.start() + child.join(timeout=5) + assert result_q.get(timeout=1) == "busy", ( + "outer lock should still be held by parent after inner re-entrant exit" + ) + + +def _try_acquire_expect_busy(palace_path, result_q): + """Helper: try to acquire, push 'busy' (raised) or 'free' (acquired) into queue.""" + try: + with mine_palace_lock(palace_path): + result_q.put("free") + except MineAlreadyRunning: + result_q.put("busy") def test_mine_global_lock_is_alias_for_back_compat(tmp_path, monkeypatch): diff --git a/tests/test_repair.py b/tests/test_repair.py index bc770ddfbe..eec9a1ed58 100644 --- a/tests/test_repair.py +++ b/tests/test_repair.py @@ -2,7 +2,7 @@ import os import sqlite3 -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest @@ -28,6 +28,16 @@ def test_get_palace_path_fallback(): assert ".mempalace" in result +def test_get_collection_name_from_config(): + from mempalace.config import get_configured_collection_name + + get_configured_collection_name.cache_clear() + with patch("mempalace.config.MempalaceConfig") as mock_config_cls: + mock_config_cls.return_value.collection_name = "custom_drawers" + assert repair._get_collection_name() == "custom_drawers" + get_configured_collection_name.cache_clear() + + # ── _paginate_ids ───────────────────────────────────────────────────── @@ -229,8 +239,11 @@ def test_rebuild_index_success(mock_backend_cls, mock_shutil, tmp_path): } mock_new_col = MagicMock() + mock_new_col.count.return_value = 2 + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 mock_backend = _install_mock_backend(mock_backend_cls, mock_col) - mock_backend.create_collection.return_value = mock_new_col + mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] repair.rebuild_index(palace_path=str(tmp_path)) @@ -239,14 +252,74 @@ def test_rebuild_index_success(mock_backend_cls, mock_shutil, tmp_path): assert "chroma.sqlite3" in str(mock_shutil.copy2.call_args) # Verify: deleted and recreated (cosine is the backend default) - mock_backend.delete_collection.assert_called_once_with(str(tmp_path), "mempalace_drawers") - mock_backend.create_collection.assert_called_once_with(str(tmp_path), "mempalace_drawers") + assert mock_backend.create_collection.call_args_list == [ + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers"), + ] + assert mock_backend.delete_collection.call_args_list == [ + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers"), + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + ] # Verify: used upsert not add + mock_temp_col.upsert.assert_called_once() mock_new_col.upsert.assert_called_once() mock_new_col.add.assert_not_called() +@patch("mempalace.repair.shutil") +@patch("mempalace.repair.ChromaBackend") +def test_rebuild_index_ignores_missing_temp_collection_at_start( + mock_backend_cls, mock_shutil, tmp_path +): + sqlite_path = tmp_path / "chroma.sqlite3" + sqlite_path.write_text("fake") + + def _fake_copy2(src, dst): + with open(dst, "w") as handle: + handle.write("backup") + + mock_shutil.copy2.side_effect = _fake_copy2 + + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + + mock_new_col = MagicMock() + mock_new_col.count.return_value = 2 + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_backend = _install_mock_backend(mock_backend_cls, mock_col) + mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] + mock_backend.delete_collection.side_effect = [ + ValueError("Collection [mempalace_drawers__repair_tmp] does not exist"), + None, + None, + ] + + repair.rebuild_index(palace_path=str(tmp_path)) + + assert mock_shutil.copy2.call_count == 1 + assert mock_backend.delete_collection.call_args_list == [ + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers"), + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + ] + + +def test_delete_collection_if_exists_reraises_unexpected_value_error(): + mock_backend = MagicMock() + mock_backend.delete_collection.side_effect = ValueError("invalid collection name") + + with pytest.raises(ValueError, match="invalid collection name"): + repair._delete_collection_if_exists(mock_backend, "/palace", "bad/name") + + @patch("mempalace.repair.shutil") @patch("mempalace.repair.ChromaBackend") def test_rebuild_index_error_reading(mock_backend_cls, mock_shutil, tmp_path): @@ -267,6 +340,21 @@ def test_check_extraction_safety_passes_when_counts_match(tmp_path): repair.check_extraction_safety(str(tmp_path), 500) +def test_check_extraction_safety_uses_configured_collection(tmp_path): + with patch("mempalace.repair.sqlite_drawer_count", return_value=500) as count: + repair.check_extraction_safety(str(tmp_path), 500, collection_name="custom_drawers") + count.assert_called_once_with(str(tmp_path), "custom_drawers") + + +def test_check_extraction_safety_default_uses_configured_collection(tmp_path): + with ( + patch("mempalace.repair._get_collection_name", return_value="custom_drawers"), + patch("mempalace.repair.sqlite_drawer_count", return_value=500) as count, + ): + repair.check_extraction_safety(str(tmp_path), 500) + count.assert_called_once_with(str(tmp_path), "custom_drawers") + + def test_check_extraction_safety_passes_when_sqlite_unreadable_and_under_cap(tmp_path): """SQLite check fails (None) but extraction is well under the cap → safe.""" with patch("mempalace.repair.sqlite_drawer_count", return_value=None): @@ -321,6 +409,73 @@ def test_sqlite_drawer_count_returns_none_on_unreadable_schema(tmp_path): assert repair.sqlite_drawer_count(str(tmp_path)) is None +@patch("mempalace.repair.shutil") +@patch("mempalace.repair.ChromaBackend") +def test_rebuild_index_default_uses_configured_collection(mock_backend_cls, mock_shutil, tmp_path): + sqlite_path = tmp_path / "chroma.sqlite3" + sqlite_path.write_text("fake") + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_new_col = MagicMock() + mock_new_col.count.return_value = 2 + mock_backend = _install_mock_backend(mock_backend_cls, mock_col) + mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] + + with ( + patch("mempalace.repair._get_collection_name", return_value="custom_drawers"), + patch("mempalace.repair.sqlite_drawer_count", return_value=2) as count, + ): + repair.rebuild_index(palace_path=str(tmp_path)) + + mock_backend.get_collection.assert_called_once_with(str(tmp_path), "custom_drawers") + count.assert_called_once_with(str(tmp_path), "custom_drawers") + assert mock_backend.create_collection.call_args_list == [ + call(str(tmp_path), "custom_drawers__repair_tmp"), + call(str(tmp_path), "custom_drawers"), + ] + assert mock_backend.delete_collection.call_args_list == [ + call(str(tmp_path), "custom_drawers__repair_tmp"), + call(str(tmp_path), "custom_drawers"), + call(str(tmp_path), "custom_drawers__repair_tmp"), + ] + + +def test_status_default_uses_configured_drawer_collection(tmp_path): + with ( + patch("mempalace.repair._get_collection_name", return_value="custom_drawers"), + patch("mempalace.repair.hnsw_capacity_status") as capacity_status, + ): + capacity_status.side_effect = [ + { + "sqlite_count": 1, + "hnsw_count": 1, + "divergence": 0, + "diverged": False, + "status": "ok", + "message": "", + }, + { + "sqlite_count": 0, + "hnsw_count": 0, + "divergence": 0, + "diverged": False, + "status": "ok", + "message": "", + }, + ] + repair.status(palace_path=str(tmp_path)) + + assert capacity_status.call_args_list[0].args == (str(tmp_path), "custom_drawers") + assert capacity_status.call_args_list[1].args == (str(tmp_path), "mempalace_closets") + + @patch("mempalace.repair.shutil") @patch("mempalace.repair.ChromaBackend") def test_rebuild_index_aborts_on_truncation_signal(mock_backend_cls, mock_shutil, tmp_path): @@ -365,19 +520,261 @@ def test_rebuild_index_proceeds_with_override(mock_backend_cls, mock_shutil, tmp }, {"ids": [], "documents": [], "metadatas": []}, ] + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 10_000 mock_new_col = MagicMock() + mock_new_col.count.return_value = 10_000 mock_backend.get_collection.return_value = mock_col - mock_backend.create_collection.return_value = mock_new_col + mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] mock_backend_cls.return_value = mock_backend with patch("mempalace.repair.sqlite_drawer_count", return_value=67_580): repair.rebuild_index(palace_path=str(tmp_path), confirm_truncation_ok=True) - mock_backend.delete_collection.assert_called_once() - mock_backend.create_collection.assert_called_once() + assert mock_backend.delete_collection.call_count == 3 + assert mock_backend.create_collection.call_count == 2 + mock_temp_col.upsert.assert_called() mock_new_col.upsert.assert_called() +@patch("mempalace.repair.shutil") +@patch("mempalace.repair.ChromaBackend") +def test_rebuild_index_stage_failure_leaves_live_collection_untouched( + mock_backend_cls, mock_shutil, tmp_path +): + sqlite_path = tmp_path / "chroma.sqlite3" + sqlite_path.write_text("fake") + + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 1 + mock_backend = _install_mock_backend(mock_backend_cls, mock_col) + mock_backend.create_collection.return_value = mock_temp_col + + with pytest.raises(repair.RebuildCollectionError) as excinfo: + repair.rebuild_index(palace_path=str(tmp_path)) + + assert excinfo.value.live_replaced is False + assert mock_shutil.copy2.call_count == 1 + assert mock_backend.delete_collection.call_args_list == [ + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + ] + + +@patch("mempalace.repair.shutil") +@patch("mempalace.repair.ChromaBackend") +def test_rebuild_index_live_failure_restores_backup(mock_backend_cls, mock_shutil, tmp_path): + sqlite_path = tmp_path / "chroma.sqlite3" + sqlite_path.write_text("fake") + + def _fake_copy2(src, dst): + with open(dst, "w") as handle: + handle.write("backup") + + mock_shutil.copy2.side_effect = _fake_copy2 + + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_new_col = MagicMock() + mock_new_col.upsert.side_effect = RuntimeError("live upsert failed") + active_backend = MagicMock() + active_backend.get_collection.return_value = mock_col + active_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] + helper_backend = MagicMock() + mock_backend_cls.side_effect = [active_backend, helper_backend] + + with pytest.raises(repair.RebuildCollectionError) as excinfo: + repair.rebuild_index(palace_path=str(tmp_path)) + + assert excinfo.value.live_replaced is True + assert mock_shutil.copy2.call_count == 2 + assert active_backend.delete_collection.call_args_list == [ + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers"), + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers"), + ] + active_backend.close_palace.assert_called_once_with(str(tmp_path)) + helper_backend.close_palace.assert_not_called() + + +@patch("mempalace.repair.shutil") +@patch("mempalace.repair.ChromaBackend") +def test_rebuild_index_live_delete_missing_still_restores_backup( + mock_backend_cls, mock_shutil, tmp_path +): + sqlite_path = tmp_path / "chroma.sqlite3" + sqlite_path.write_text("fake") + + def _fake_copy2(src, dst): + with open(dst, "w") as handle: + handle.write("backup") + + mock_shutil.copy2.side_effect = _fake_copy2 + + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_backend = _install_mock_backend(mock_backend_cls, mock_col) + mock_backend.create_collection.side_effect = [mock_temp_col, RuntimeError("create failed")] + mock_backend.delete_collection.side_effect = [ + None, + None, + None, + repair.ChromaNotFoundError("missing"), + ] + + with pytest.raises(repair.RebuildCollectionError) as excinfo: + repair.rebuild_index(palace_path=str(tmp_path)) + + assert excinfo.value.live_replaced is True + assert mock_shutil.copy2.call_count == 2 + assert mock_backend.delete_collection.call_args_list == [ + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers"), + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers"), + ] + + +@patch("mempalace.repair.shutil") +@patch("mempalace.repair.ChromaBackend") +def test_rebuild_index_restore_failure_preserves_original_error( + mock_backend_cls, mock_shutil, tmp_path, capsys +): + sqlite_path = tmp_path / "chroma.sqlite3" + sqlite_path.write_text("fake") + + def _copy2_side_effect(src, dst): + if str(src).endswith(".backup"): + raise PermissionError("locked sqlite") + with open(dst, "w") as handle: + handle.write("backup") + + mock_shutil.copy2.side_effect = _copy2_side_effect + + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_new_col = MagicMock() + mock_new_col.upsert.side_effect = RuntimeError("live upsert failed") + mock_backend = _install_mock_backend(mock_backend_cls, mock_col) + mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] + + with pytest.raises(repair.RebuildCollectionError) as excinfo: + repair.rebuild_index(palace_path=str(tmp_path)) + + out = capsys.readouterr().out + assert "locked sqlite" in out + assert "Manual restore required" in out + assert "live upsert failed" in str(excinfo.value) + + +@patch("mempalace.repair.ChromaBackend") +def test_rebuild_collection_via_temp_keeps_original_error_when_cleanup_fails( + mock_backend_cls, +): + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_backend = _install_mock_backend(mock_backend_cls, mock_col) + mock_backend.create_collection.side_effect = [mock_temp_col, RuntimeError("live build failed")] + mock_backend.delete_collection.side_effect = [ + None, + None, + RuntimeError("cleanup failed"), + ] + + with pytest.raises(repair.RebuildCollectionError) as excinfo: + repair._rebuild_collection_via_temp( + mock_backend, + "/palace", + ["id1", "id2"], + ["doc1", "doc2"], + [{"wing": "a"}, {"wing": "b"}], + batch_size=5000, + progress=lambda *args, **kwargs: None, + ) + + assert "live build failed" in str(excinfo.value) + assert excinfo.value.live_replaced is True + assert mock_backend.delete_collection.call_args_list == [ + call("/palace", "mempalace_drawers__repair_tmp"), + call("/palace", "mempalace_drawers"), + call("/palace", "mempalace_drawers__repair_tmp"), + ] + + +@patch("mempalace.repair.shutil") +@patch("mempalace.repair.ChromaBackend") +def test_rebuild_index_ignores_temp_cleanup_failure_after_success( + mock_backend_cls, mock_shutil, tmp_path +): + sqlite_path = tmp_path / "chroma.sqlite3" + sqlite_path.write_text("fake") + + def _fake_copy2(src, dst): + with open(dst, "w") as handle: + handle.write("backup") + + mock_shutil.copy2.side_effect = _fake_copy2 + + mock_col = MagicMock() + mock_col.count.return_value = 2 + mock_col.get.return_value = { + "ids": ["id1", "id2"], + "documents": ["doc1", "doc2"], + "metadatas": [{"wing": "a"}, {"wing": "b"}], + } + mock_temp_col = MagicMock() + mock_temp_col.count.return_value = 2 + mock_new_col = MagicMock() + mock_new_col.count.return_value = 2 + mock_backend = _install_mock_backend(mock_backend_cls, mock_col) + mock_backend.create_collection.side_effect = [mock_temp_col, mock_new_col] + mock_backend.delete_collection.side_effect = [ + None, + None, + RuntimeError("cleanup failed"), + ] + + repair.rebuild_index(palace_path=str(tmp_path)) + + assert mock_shutil.copy2.call_count == 1 + assert mock_backend.delete_collection.call_args_list == [ + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + call(str(tmp_path), "mempalace_drawers"), + call(str(tmp_path), "mempalace_drawers__repair_tmp"), + ] + + # ── repair_max_seq_id ───────────────────────────────────────────────── diff --git a/tests/test_searcher.py b/tests/test_searcher.py index 6a32fcd940..6e3c676fcf 100644 --- a/tests/test_searcher.py +++ b/tests/test_searcher.py @@ -84,6 +84,24 @@ def test_search_memories_query_error(self): assert "error" in result assert "query failed" in result["error"] + def test_search_memories_vector_path_uses_explicit_collection_name(self): + mock_col = MagicMock() + mock_col.query.return_value = { + "documents": [[]], + "metadatas": [[]], + "distances": [[]], + "ids": [[]], + } + + with patch("mempalace.searcher.get_collection", return_value=mock_col) as get_collection: + search_memories("test", "/fake/path", collection_name="custom_drawers") + + get_collection.assert_called_once_with( + "/fake/path", + collection_name="custom_drawers", + create=False, + ) + def test_search_memories_filters_in_result(self, palace_path, seeded_collection): result = search_memories("test", palace_path, wing="project", room="backend") assert result["filters"]["wing"] == "project" @@ -102,7 +120,7 @@ def test_search_memories_handles_none_metadata(self): "ids": [["d1", "d2"]], } - def mock_get_collection(path, create=False): + def mock_get_collection(path, collection_name=None, create=False): # First call: drawers. Second call: closets — raise so hybrid # degrades to pure drawer search (the catch block covers it). if not hasattr(mock_get_collection, "_called"): @@ -250,9 +268,9 @@ def test_search_applies_bm25_hybrid_rerank(self, capsys): captured = capsys.readouterr() first_block, _, _ = captured.out.partition("[2]") # Lexical match must rank first - assert ( - "b.md" in first_block - ), f"expected lexical match 'b.md' at rank 1, got:\n{captured.out}" + assert "b.md" in first_block, ( + f"expected lexical match 'b.md' at rank 1, got:\n{captured.out}" + ) # Non-zero bm25 reported assert "bm25=" in first_block assert "bm25=0.0" not in first_block