core/patch_memory.py

"""
core/patch_memory.py — v4.9 Task Memory: append-friendly vector store.

Purpose
-------
When ``AutoDebugLoop`` fails to fix an error after N attempts, the failing
patch is recorded here. On the next run, semantically similar errors can be
matched and surfaced as "known-bad" advisories.

Design constraints
------------------
* **No external dependencies.** Pure-Python TF-IDF cosine similarity.
* **Bounded size.** ``MAX_ENTRIES`` keeps the in-memory corpus small.
* **Append-friendly persistence.** New failures are appended to a JSONL log
  and compacted back into ``patch_memory.json`` periodically. This avoids
  rewriting the full snapshot file on every failed attempt.
* **Defensive-only.** The memory is advisory data, never executable code.
"""

from __future__ import annotations

import json
import math
import os
import re
import threading
import time
from datetime import datetime, timezone
from typing import Any, Dict, Iterable, List, Optional, Tuple

try:
    from core.workspace_state import ensure_persistent_workspace  # noqa: WPS433
except Exception:  # pragma: no cover
    def ensure_persistent_workspace(path: Optional[str] = None) -> str:  # type: ignore[misc]
        root = os.path.abspath(
            path or os.environ.get(
                "DEVAI_PERSISTENT_WORKSPACE_DIR",
                os.path.join(os.environ.get("WORKSPACE_DIR", "/tmp/devai_workspace"), "persistent"),
            )
        )
        os.makedirs(root, exist_ok=True)
        return root


MAX_ENTRIES = 500
DEFAULT_SIMILARITY_THRESHOLD = 0.55
DEFAULT_TOP_K = 3
PATCH_SNIPPET_CHARS = 1200
ERROR_SNIPPET_CHARS = 800
COMPACT_EVERY = 25

_STOP_TOKENS = {
    "the", "a", "an", "of", "to", "in", "on", "for", "and", "or", "is",
    "was", "be", "at", "by", "with", "as", "this", "that", "it", "if",
    "not", "no", "yes", "from", "into", "then", "self", "cls", "def",
    "return", "class", "import", "true", "false", "none",
    "file", "line", "of",
}
_TOKEN_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]{1,}")


def _tokenize(text: str) -> List[str]:
    if not text:
        return []
    out: List[str] = []
    for match in _TOKEN_RE.finditer(text):
        token = match.group(0).lower()
        if len(token) < 3 or token in _STOP_TOKENS:
            continue
        out.append(token)
    return out


def _term_freq(tokens: Iterable[str]) -> Dict[str, float]:
    counts: Dict[str, float] = {}
    for token in tokens:
        counts[token] = counts.get(token, 0.0) + 1.0
    if not counts:
        return {}
    return {key: 1.0 + math.log(val) for key, val in counts.items()}


def _norm(vec: Dict[str, float]) -> float:
    return math.sqrt(sum(v * v for v in vec.values())) or 1.0


def _cosine(a: Dict[str, float], a_norm: float, b: Dict[str, float], b_norm: float) -> float:
    if not a or not b:
        return 0.0
    if len(a) > len(b):
        a, b = b, a
        a_norm, b_norm = b_norm, a_norm
    dot = 0.0
    for key, val in a.items():
        other = b.get(key)
        if other:
            dot += val * other
    return dot / (a_norm * b_norm)


def _build_query_text(diagnostic: Dict[str, Any]) -> str:
    parts = [
        str(diagnostic.get("error_type") or ""),
        str(diagnostic.get("message") or ""),
        str(diagnostic.get("file_path") or ""),
        str(diagnostic.get("raw_log") or "")[:ERROR_SNIPPET_CHARS],
    ]
    return "\n".join(part for part in parts if part)


class PatchMemory:
    """Thread-safe append-friendly vector store of failed patch attempts."""

    def __init__(self, workspace: Optional[str] = None, max_entries: int = MAX_ENTRIES) -> None:
        self._workspace = ensure_persistent_workspace(workspace)
        self._path = os.path.join(self._workspace, "patch_memory.json")
        self._log_path = os.path.join(self._workspace, "patch_memory.jsonl")
        self._max_entries = int(max_entries)
        self._lock = threading.RLock()
        self._data: Dict[str, Any] = {"version": 2, "entries": []}
        self._dirty_events = 0
        self._load()

    def _load_snapshot(self) -> Dict[str, Any]:
        if not os.path.exists(self._path):
            return {"version": 2, "entries": []}
        try:
            with open(self._path, "r", encoding="utf-8") as fh:
                payload = json.load(fh)
            if isinstance(payload, dict) and isinstance(payload.get("entries"), list):
                for entry in payload["entries"]:
                    vec = entry.get("vector") or {}
                    if "norm" not in entry:
                        entry["norm"] = _norm(vec)
                payload["version"] = 2
                return payload
        except Exception:
            try:
                os.replace(self._path, self._path + ".corrupt")
            except OSError:
                pass
        return {"version": 2, "entries": []}

    def _load(self) -> None:
        self._data = self._load_snapshot()
        if not os.path.exists(self._log_path):
            return
        try:
            with open(self._log_path, "r", encoding="utf-8") as fh:
                for line in fh:
                    line = line.strip()
                    if not line:
                        continue
                    payload = json.loads(line)
                    if not isinstance(payload, dict):
                        continue
                    entry = payload.get("entry") if payload.get("op") == "append" else None
                    if isinstance(entry, dict):
                        self._append_entry(entry, compacting=False)
                    elif payload.get("op") == "clear":
                        self._data = {"version": 2, "entries": []}
        except Exception:
            try:
                os.replace(self._log_path, self._log_path + ".corrupt")
            except OSError:
                pass
            self._data = self._load_snapshot()
        self._dirty_events = 0

    def _save_snapshot(self) -> None:
        os.makedirs(os.path.dirname(self._path), exist_ok=True)
        tmp = self._path + ".tmp"
        with open(tmp, "w", encoding="utf-8") as fh:
            json.dump(self._data, fh, ensure_ascii=False, indent=2)
        os.replace(tmp, self._path)

    def _append_log(self, payload: Dict[str, Any]) -> None:
        os.makedirs(os.path.dirname(self._log_path), exist_ok=True)
        with open(self._log_path, "a", encoding="utf-8") as fh:
            fh.write(json.dumps(payload, ensure_ascii=False, default=str) + "\n")

    def _compact(self) -> None:
        self._save_snapshot()
        tmp = self._log_path + ".tmp"
        with open(tmp, "w", encoding="utf-8") as fh:
            fh.write("")
        os.replace(tmp, self._log_path)
        self._dirty_events = 0

    def _append_entry(self, entry: Dict[str, Any], compacting: bool) -> None:
        entries: List[Dict[str, Any]] = self._data.setdefault("entries", [])
        entries.append(entry)
        if len(entries) > self._max_entries:
            overflow = len(entries) - self._max_entries
            del entries[:overflow]
        if compacting:
            self._dirty_events += 1
            if self._dirty_events >= COMPACT_EVERY or len(entries) >= self._max_entries:
                self._compact()

    @property
    def path(self) -> str:
        return self._path

    @property
    def log_path(self) -> str:
        return self._log_path

    def __len__(self) -> int:
        with self._lock:
            return len(self._data.get("entries", []))

    def record_failure(
        self,
        diagnostic: Dict[str, Any],
        patch_text: str,
        *,
        attempts: int = 0,
        tags: Optional[List[str]] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        query_text = _build_query_text(diagnostic)
        tokens = _tokenize(query_text + "\n" + (patch_text or ""))
        tf = _term_freq(tokens)
        entry = {
            "id": f"fp-{int(time.time() * 1000)}-{len(self._data.get('entries', [])) + 1}",
            "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "error_type": str(diagnostic.get("error_type") or "Unknown"),
            "error_message": str(diagnostic.get("message") or "")[:ERROR_SNIPPET_CHARS],
            "file_path": str(diagnostic.get("file_path") or ""),
            "line_number": int(diagnostic.get("line_number") or 0),
            "attempts": int(attempts),
            "patch_snippet": (patch_text or "")[:PATCH_SNIPPET_CHARS],
            "tags": list(tags or []),
            "extra": dict(extra or {}),
            "vector": tf,
            "norm": _norm(tf),
        }
        with self._lock:
            self._append_entry(entry, compacting=False)
            self._append_log({"op": "append", "entry": entry})
            self._dirty_events += 1
            if self._dirty_events >= COMPACT_EVERY or len(self._data.get("entries", [])) >= self._max_entries:
                self._compact()
        return {key: val for key, val in entry.items() if key not in ("vector", "norm")}

    def query_similar(
        self,
        diagnostic: Dict[str, Any],
        *,
        top_k: int = DEFAULT_TOP_K,
        threshold: float = DEFAULT_SIMILARITY_THRESHOLD,
    ) -> List[Dict[str, Any]]:
        query_text = _build_query_text(diagnostic)
        q_tokens = _tokenize(query_text)
        if not q_tokens:
            return []
        q_vec = _term_freq(q_tokens)
        q_norm = _norm(q_vec)
        target_file = str(diagnostic.get("file_path") or "").strip()
        target_error = str(diagnostic.get("error_type") or "").strip()
        with self._lock:
            candidates: List[Tuple[float, Dict[str, Any]]] = []
            for entry in self._data.get("entries", []):
                vec = entry.get("vector") or {}
                e_norm = float(entry.get("norm") or 1.0)
                sim = _cosine(q_vec, q_norm, vec, e_norm)
                if target_error and entry.get("error_type") == target_error:
                    sim += 0.10
                if target_file and entry.get("file_path") == target_file:
                    sim += 0.05
                if sim >= threshold:
                    candidates.append((sim, entry))
            candidates.sort(key=lambda item: item[0], reverse=True)
            results: List[Dict[str, Any]] = []
            for sim, entry in candidates[:top_k]:
                results.append({
                    "id": entry.get("id"),
                    "similarity": round(min(sim, 1.0), 4),
                    "error_type": entry.get("error_type"),
                    "error_message": entry.get("error_message"),
                    "file_path": entry.get("file_path"),
                    "line_number": entry.get("line_number"),
                    "attempts": entry.get("attempts"),
                    "patch_snippet": entry.get("patch_snippet"),
                    "created_at": entry.get("created_at"),
                    "tags": entry.get("tags", []),
                })
            return results

    def format_advisory(self, matches: List[Dict[str, Any]]) -> str:
        if not matches:
            return ""
        lines = [
            "── PATCH MEMORY: previously-failed fixes for similar errors ──",
            "The following patches were tried before for a semantically similar",
            "error and DID NOT WORK. Do NOT repeat these patterns:",
            "",
        ]
        for idx, match in enumerate(matches, 1):
            lines.append(
                f"[{idx}] {match['error_type']} @ {match['file_path']}:{match['line_number']} "
                f"(similarity={match['similarity']:.2f}, {match['attempts']} attempt(s))"
            )
            if match.get("error_message"):
                lines.append(f"    error: {match['error_message']}")
            snippet = (match.get("patch_snippet") or "").strip()
            if snippet:
                lines.append("    failed-patch-snippet:")
                for row in snippet.splitlines()[:20]:
                    lines.append("      | " + row)
            lines.append("")
        lines.append("── END PATCH MEMORY ──")
        return "\n".join(lines)

    def clear(self) -> None:
        with self._lock:
            self._data = {"version": 2, "entries": []}
            self._append_log({"op": "clear"})
            self._compact()

    def stats(self) -> Dict[str, Any]:
        with self._lock:
            entries = self._data.get("entries", [])
            error_types: Dict[str, int] = {}
            for entry in entries:
                error_type = entry.get("error_type") or "Unknown"
                error_types[error_type] = error_types.get(error_type, 0) + 1
            return {
                "path": self._path,
                "log_path": self._log_path,
                "count": len(entries),
                "max_entries": self._max_entries,
                "dirty_events": self._dirty_events,
                "error_types": dict(sorted(error_types.items(), key=lambda kv: -kv[1])[:20]),
            }


_default_memory: Optional[PatchMemory] = None
_default_lock = threading.Lock()


def get_default_memory(workspace: Optional[str] = None) -> PatchMemory:
    global _default_memory
    with _default_lock:
        if _default_memory is None:
            _default_memory = PatchMemory(workspace=workspace)
        return _default_memory


def reset_default_memory() -> None:
    global _default_memory
    with _default_lock:
        _default_memory = None


__all__ = [
    "PatchMemory",
    "get_default_memory",
    "reset_default_memory",
    "MAX_ENTRIES",
    "DEFAULT_SIMILARITY_THRESHOLD",
    "DEFAULT_TOP_K",
]