""" core/zip_handler.py — Zip Intelligence helpers (v4.5 — Line-Locked Edition). Features -------- - safely extract uploaded ZIP archives into an isolated workspace - inspect file structure and text-file contents for AI context building - apply AI-suggested changes back onto the extracted project using THREE strategies (in strict priority order): 1. **Line-Locked Patch blocks** (```linepatch``` fences) — surgical edits bounded by an explicit ``LINES: -`` header. Any content the model produces outside that range is REJECTED. This is the preferred format for Phase-1 reliability. 2. **Unified Diff / Git patch** blocks — applied via ``git apply``. 3. **# FILE: full-file rewrite** blocks — kept for backwards compatibility but now flagged with ``strategy="file_blocks"`` so the pipeline can downgrade its trust score. - recompress the modified project into a downloadable ZIP. """ from __future__ import annotations import os import re import shutil import subprocess import tempfile import zipfile from pathlib import Path from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple IGNORE_DIRS = { ".git", ".venv", "venv", "node_modules", "__pycache__", ".idea", ".vscode", "dist", "build", ".next", ".cache", "__MACOSX", } TEXT_EXTENSIONS = { ".py", ".js", ".ts", ".tsx", ".jsx", ".java", ".go", ".rs", ".c", ".cpp", ".h", ".hpp", ".rb", ".php", ".sh", ".yml", ".yaml", ".json", ".md", ".html", ".css", ".sql", ".toml", ".cfg", ".ini", ".txt", ".env", ".dockerfile", ".xml", } _CODE_BLOCK_RE = re.compile(r"```(?P[\w.+-]*)\n(?P.*?)```", re.S) _FILE_HEADER_RE = re.compile(r"^[#/;!\- ]*FILE:\s*([^\n\r]+)", re.I | re.M) _DIFF_PLUS_RE = re.compile(r"^\+\+\+\s+(?:b/)?(.+)$", re.M) # --------------------------------------------------------------------------- # Line-Locked patch format # --------------------------------------------------------------------------- # Expected block (fence language: ``linepatch``, ``line-patch``, or ``patch-lines``): # # ```linepatch # FILE: relative/path/to/file.py # LINES: 42-57 # --- # # ``` # # Multiple linepatch blocks may appear in one response, even against the same # file. They are applied bottom-up (highest starting line first) so earlier # line numbers remain valid. _LINEPATCH_LANGS = {"linepatch", "line-patch", "patch-lines", "surgical"} _LP_FILE_RE = re.compile(r"^\s*FILE:\s*(?P[^\n\r]+)\s*$", re.I | re.M) _LP_LINES_RE = re.compile( r"^\s*LINES:\s*(?P\d+)\s*-\s*(?P\d+)\s*$", re.I | re.M, ) _LP_SEPARATOR_RE = re.compile(r"^\s*---+\s*$", re.M) def _safe_rel_path(rel_path: str) -> str: rel = str(rel_path or "").replace("\\", "/").strip().lstrip("/") rel = re.sub(r"/+", "/", rel) if not rel or rel in {".", ".."}: raise ValueError("empty relative path") parts = [p for p in rel.split("/") if p not in {"", "."}] if any(part == ".." for part in parts): raise ValueError(f"unsafe path: {rel_path}") return "/".join(parts) def _collect_candidate_modified_files(text: str) -> List[str]: candidates: List[str] = [] for block in _extract_linepatch_blocks(text): path = block.get("path") if path and path not in candidates: candidates.append(path) for block in _extract_diff_blocks(text): for path in _DIFF_PLUS_RE.findall(block): raw = (path or "").strip() if not raw or raw == "/dev/null": continue try: clean = _safe_rel_path(raw) except ValueError: continue if clean not in candidates: candidates.append(clean) for path, _content in _extract_file_blocks(text): if path not in candidates: candidates.append(path) return candidates def _derive_allowed_dirs(allowed_files: Sequence[str]) -> List[str]: dirs = set() for rel in allowed_files: clean = _safe_rel_path(rel) parent = clean.rsplit("/", 1)[0] if "/" in clean else "" if parent: dirs.add(parent) return sorted(dirs) def _enforce_diff_guard(text: str, allowed_files: Optional[Sequence[str]] = None) -> Optional[Dict[str, Any]]: if not allowed_files: return None normalized_files = [] for rel in allowed_files: try: clean = _safe_rel_path(rel) except ValueError: continue if clean not in normalized_files: normalized_files.append(clean) if not normalized_files: return None allowed_set = set(normalized_files) allowed_dirs = _derive_allowed_dirs(normalized_files) touched = _collect_candidate_modified_files(text) blocked: List[str] = [] for rel in touched: if rel in allowed_set: continue if any(rel.startswith(directory + "/") for directory in allowed_dirs): continue blocked.append(rel) if not blocked: return None return { "applied": False, "strategy": "diff_guard", "modified_files": [], "errors": [ "Diff Guard rejected patch outside GraphMapper scope: " + ", ".join(blocked) ], "blocked_files": blocked, "allowed_scope": normalized_files, } def _looks_like_text(path: Path, sample_size: int = 2048) -> bool: ext = path.suffix.lower() if ext in TEXT_EXTENSIONS: return True try: with open(path, "rb") as fh: sample = fh.read(sample_size) if b"\x00" in sample: return False sample.decode("utf-8") return True except Exception: return False def _iter_project_files(project_root: str) -> Iterable[Tuple[str, Path]]: root = Path(project_root).resolve() for dirpath, dirnames, filenames in os.walk(root): dirnames[:] = [d for d in dirnames if d not in IGNORE_DIRS] for filename in sorted(filenames): abs_path = Path(dirpath) / filename rel_path = abs_path.relative_to(root).as_posix() yield rel_path, abs_path def _detect_archive_root(zf: zipfile.ZipFile) -> str: top_levels = set() for name in zf.namelist(): clean = name.replace("\\", "/").strip("/") if not clean or clean.startswith("__MACOSX/"): continue parts = clean.split("/") if parts and parts[0]: top_levels.add(parts[0]) if len(top_levels) > 1: return "" return next(iter(top_levels), "") if len(top_levels) == 1 else "" def extract_uploaded_zip( zip_path: str, workspace: Optional[str] = None, project_id: Optional[str] = None, ) -> Dict[str, Any]: """Extract an uploaded ZIP into an isolated temp directory. Returns metadata needed for later previewing and re-zipping. """ if not zip_path or not os.path.exists(zip_path): raise FileNotFoundError(f"zip not found: {zip_path}") base_workspace = workspace or os.path.join(tempfile.gettempdir(), "devai_zipintelligence") os.makedirs(base_workspace, exist_ok=True) prefix = f"{project_id}_" if project_id else "zipintel_" extraction_dir = tempfile.mkdtemp(prefix=prefix, dir=base_workspace) with zipfile.ZipFile(zip_path, "r") as zf: archive_root_prefix = _detect_archive_root(zf) zf.extractall(extraction_dir) project_root = os.path.join(extraction_dir, archive_root_prefix) if archive_root_prefix else extraction_dir project_root = os.path.abspath(project_root) if not os.path.isdir(project_root): raise FileNotFoundError("ZIP extracted, but no project root directory was found") return { "zip_path": os.path.abspath(zip_path), "zip_name": os.path.basename(zip_path), "extract_dir": os.path.abspath(extraction_dir), "project_root": project_root, "archive_root_prefix": archive_root_prefix, } def read_project_file(project_root: str, relative_path: str, max_chars: int = 12000) -> str: rel = _safe_rel_path(relative_path) abs_path = (Path(project_root).resolve() / rel).resolve() root = Path(project_root).resolve() if root not in abs_path.parents and abs_path != root: raise ValueError(f"path escapes project root: {relative_path}") if not abs_path.exists() or not abs_path.is_file(): raise FileNotFoundError(relative_path) if not _looks_like_text(abs_path): return "[binary or non-text file omitted]" with open(abs_path, "r", encoding="utf-8", errors="ignore") as fh: text = fh.read(max_chars + 1) if len(text) > max_chars: return text[:max_chars] + "\n...[truncated]" return text def list_project_files(project_root: str, max_files: int = 1000) -> List[str]: files: List[str] = [] for rel_path, _ in _iter_project_files(project_root): files.append(rel_path) if len(files) >= max_files: break return files def build_file_tree(project_root: str, max_entries: int = 250) -> str: lines = ["."] count = 0 root = Path(project_root).resolve() for dirpath, dirnames, filenames in os.walk(root): dirnames[:] = sorted([d for d in dirnames if d not in IGNORE_DIRS]) rel_dir = Path(dirpath).resolve().relative_to(root).as_posix() depth = 0 if rel_dir == "." else rel_dir.count("/") + 1 base_indent = " " * depth if rel_dir != ".": lines.append(f"{base_indent}{Path(dirpath).name}/") count += 1 if count >= max_entries: lines.append(" ...[tree truncated]") break for filename in sorted(filenames): rel_file = (Path(dirpath) / filename).resolve().relative_to(root).as_posix() lines.append(f"{base_indent} {filename}") count += 1 if count >= max_entries: lines.append(" ...[tree truncated]") return "\n".join(lines) return "\n".join(lines) def build_zip_ai_context( project_root: str, focus_files: Optional[Sequence[str]] = None, max_files: int = 18, per_file_chars: int = 8000, max_total_chars: int = 70000, include_line_numbers: bool = True, ) -> str: """Build the context blob shown to the AI. When ``include_line_numbers`` is True (the v4.5 default), each source file is rendered with a ``NNN: `` prefix on every line so the model can address edits by exact line numbers — the foundation of Line-Locking. """ files = list_project_files(project_root, max_files=1200) ordered: List[str] = [] seen = set() for rel in focus_files or []: if rel in files and rel not in seen: ordered.append(rel) seen.add(rel) for rel in files: if rel not in seen: ordered.append(rel) seen.add(rel) sections = [ "## ZIP Intelligence — Project file tree", "```text", build_file_tree(project_root), "```", "", "## ZIP Intelligence — Readable file contents (with line numbers)" if include_line_numbers else "## ZIP Intelligence — Readable file contents", ] total = sum(len(s) for s in sections) file_count = 0 for rel in ordered: abs_path = Path(project_root) / rel if not abs_path.is_file() or not _looks_like_text(abs_path): continue snippet = read_project_file(project_root, rel, max_chars=per_file_chars) if include_line_numbers: numbered_lines = [] for idx, line in enumerate(snippet.splitlines(), start=1): numbered_lines.append(f"{idx:>4}: {line}") snippet = "\n".join(numbered_lines) block = f"### FILE: {rel}\n```\n{snippet}\n```\n" if total + len(block) > max_total_chars and file_count > 0: sections.append("\n...[additional files omitted to stay within context budget]") break sections.append(block) total += len(block) file_count += 1 if file_count >= max_files: sections.append("\n...[file limit reached]\n") break if file_count == 0: sections.append("_(No readable text files found in the uploaded ZIP.)_") return "\n".join(sections).strip() def preview_project_contents( project_root: str, max_files: int = 30, per_file_chars: int = 3000, max_total_chars: int = 90000, ) -> str: return build_zip_ai_context( project_root=project_root, focus_files=None, max_files=max_files, per_file_chars=per_file_chars, max_total_chars=max_total_chars, include_line_numbers=False, ) # --------------------------------------------------------------------------- # Strategy 1 — Line-Locked patch blocks (surgical edit) # --------------------------------------------------------------------------- def _extract_linepatch_blocks(text: str) -> List[Dict[str, Any]]: """Return a list of parsed linepatch blocks. Each entry: ``{"path": str, "start": int, "end": int, "content": str}``. Malformed blocks are silently skipped; callers can compare the number of returned entries against the number of ``linepatch`` fences to detect invalid submissions. """ blocks: List[Dict[str, Any]] = [] for match in _CODE_BLOCK_RE.finditer(text or ""): lang = (match.group("lang") or "").strip().lower() if lang not in _LINEPATCH_LANGS: continue body = match.group("body") or "" file_m = _LP_FILE_RE.search(body) lines_m = _LP_LINES_RE.search(body) if not file_m or not lines_m: continue try: path = _safe_rel_path(file_m.group("path").strip()) except ValueError: continue start = int(lines_m.group("start")) end = int(lines_m.group("end")) if start < 1 or end < start: continue # Find where the header ends. Prefer an explicit `---` separator; if # missing, take everything after the last header line. header_end = max(file_m.end(), lines_m.end()) sep_m = _LP_SEPARATOR_RE.search(body, header_end) content_start = sep_m.end() if sep_m else header_end # Skip a single leading newline so the payload starts cleanly. content = body[content_start:] if content.startswith("\n"): content = content[1:] # Strip trailing whitespace-only newline; keep meaningful indentation. content = content.rstrip("\n") blocks.append({ "path": path, "start": start, "end": end, "content": content, }) return blocks def _apply_linepatch_blocks(project_root: str, text: str) -> Dict[str, Any]: """Apply Line-Locked patches to real files on disk. * Refuses to touch lines outside the declared ``LINES:`` range. * If the target file has fewer lines than ``end``, the patch is REJECTED for that block (returned in ``errors``). * Multiple blocks against the same file are applied bottom-up so earlier line numbers stay valid after each splice. """ parsed = _extract_linepatch_blocks(text) if not parsed: return { "applied": False, "strategy": "line_locked", "modified_files": [], "patch_details": [], "errors": ["no linepatch blocks found"], } root = Path(project_root).resolve() # Group by file so we can sort bottom-up per file. per_file: Dict[str, List[Dict[str, Any]]] = {} for blk in parsed: per_file.setdefault(blk["path"], []).append(blk) modified_files: List[str] = [] patch_details: List[Dict[str, Any]] = [] errors: List[str] = [] for rel_path, blocks in per_file.items(): try: abs_path = (root / rel_path).resolve() except Exception as exc: errors.append(f"{rel_path}: bad path ({exc})") continue if root not in abs_path.parents and abs_path != root: errors.append(f"{rel_path}: escapes project root") continue if not abs_path.exists() or not abs_path.is_file(): errors.append(f"{rel_path}: file not found — line-locked patches " f"can only edit existing files") continue try: original_text = abs_path.read_text(encoding="utf-8", errors="ignore") except Exception as exc: errors.append(f"{rel_path}: read failed ({exc})") continue original_lines = original_text.splitlines(keepends=False) total_lines = len(original_lines) # Validate every block against the CURRENT (unmodified) line count. valid_blocks: List[Dict[str, Any]] = [] for blk in blocks: if blk["start"] > total_lines or blk["end"] > total_lines: errors.append( f"{rel_path}: LINES {blk['start']}-{blk['end']} out of " f"range (file has {total_lines} lines) — patch rejected" ) continue valid_blocks.append(blk) if not valid_blocks: continue # Detect overlapping ranges within the same file — reject the file # entirely to preserve the surgical-editing guarantee. sorted_by_start = sorted(valid_blocks, key=lambda b: b["start"]) for i in range(1, len(sorted_by_start)): prev = sorted_by_start[i - 1] cur = sorted_by_start[i] if cur["start"] <= prev["end"]: errors.append( f"{rel_path}: overlapping linepatch ranges " f"({prev['start']}-{prev['end']} vs " f"{cur['start']}-{cur['end']}) — file skipped" ) valid_blocks = [] break if not valid_blocks: continue # Apply bottom-up so line indices for earlier blocks stay valid. new_lines = list(original_lines) applied_ranges: List[Tuple[int, int, int]] = [] # (start, end, new_len) for blk in sorted(valid_blocks, key=lambda b: b["start"], reverse=True): start_idx = blk["start"] - 1 # inclusive, 0-based end_idx = blk["end"] # exclusive slice bound replacement = blk["content"].splitlines(keepends=False) new_lines[start_idx:end_idx] = replacement applied_ranges.append((blk["start"], blk["end"], len(replacement))) # Preserve trailing newline behaviour of the original file. new_text = "\n".join(new_lines) if original_text.endswith("\n") and not new_text.endswith("\n"): new_text += "\n" try: abs_path.write_text(new_text, encoding="utf-8") except Exception as exc: errors.append(f"{rel_path}: write failed ({exc})") continue modified_files.append(rel_path) for orig_start, orig_end, new_len in sorted(applied_ranges): patch_details.append({ "file": rel_path, "original_range": [orig_start, orig_end], "original_line_count": orig_end - orig_start + 1, "replacement_line_count": new_len, }) return { "applied": bool(modified_files), "strategy": "line_locked", "modified_files": modified_files, "patch_details": patch_details, "errors": errors, } # --------------------------------------------------------------------------- # Strategy 2 — Unified diff / git patch # --------------------------------------------------------------------------- def _extract_diff_blocks(text: str) -> List[str]: blocks: List[str] = [] for match in _CODE_BLOCK_RE.finditer(text or ""): lang = (match.group("lang") or "").strip().lower() body = match.group("body") or "" if lang in {"diff", "patch"} and body.strip(): blocks.append(body) if not blocks and ("diff --git" in (text or "") or ("@@" in (text or "") and "+++" in (text or ""))): blocks.append(text) return blocks def _apply_diff_blocks(project_root: str, text: str) -> Dict[str, Any]: diff_blocks = _extract_diff_blocks(text) if not diff_blocks: return { "applied": False, "strategy": "unified_diff", "modified_files": [], "errors": ["no unified diff found"], } if not shutil.which("git"): return { "applied": False, "strategy": "unified_diff", "modified_files": [], "errors": ["`git` binary unavailable — cannot apply unified diff"], } patch_text = "\n\n".join(block.strip() for block in diff_blocks if block.strip()).strip() + "\n" patch_file = tempfile.NamedTemporaryFile("w", suffix=".patch", delete=False, encoding="utf-8") try: patch_file.write(patch_text) patch_file.close() proc = subprocess.run( ["git", "-C", project_root, "apply", "--whitespace=nowarn", "--reject", patch_file.name], capture_output=True, text=True, ) if proc.returncode != 0: return { "applied": False, "strategy": "unified_diff", "modified_files": [], "errors": [proc.stderr.strip() or proc.stdout.strip() or "git apply failed"], } modified = [ _safe_rel_path(path.strip()) for path in _DIFF_PLUS_RE.findall(patch_text) if path.strip() and path.strip() != "/dev/null" ] deduped = list(dict.fromkeys(modified)) return { "applied": True, "strategy": "unified_diff", "modified_files": deduped, "errors": [], } finally: try: os.unlink(patch_file.name) except OSError: pass # --------------------------------------------------------------------------- # Strategy 3 — # FILE: full-file rewrite (fallback / new-file support) # --------------------------------------------------------------------------- def _extract_file_blocks(text: str) -> List[Tuple[str, str]]: files: List[Tuple[str, str]] = [] for match in _CODE_BLOCK_RE.finditer(text or ""): lang = (match.group("lang") or "").strip().lower() # Skip specialised fences so a linepatch/diff never doubles as a file block. if lang in _LINEPATCH_LANGS or lang in {"diff", "patch"}: continue body = match.group("body") or "" header = _FILE_HEADER_RE.search(body) if not header: continue try: path = _safe_rel_path(header.group(1).strip()) except ValueError: continue content = body[header.end():].lstrip("\n") files.append((path, content)) return files def _apply_file_blocks(project_root: str, text: str) -> Dict[str, Any]: parsed = _extract_file_blocks(text) modified_files: List[str] = [] if not parsed: return { "applied": False, "strategy": "file_blocks", "modified_files": modified_files, "errors": ["no # FILE blocks found"], } root = Path(project_root).resolve() for rel_path, content in parsed: abs_path = (root / rel_path).resolve() if root not in abs_path.parents and abs_path != root: raise ValueError(f"unsafe write target: {rel_path}") abs_path.parent.mkdir(parents=True, exist_ok=True) with open(abs_path, "w", encoding="utf-8", errors="ignore") as fh: fh.write(content) modified_files.append(rel_path) return { "applied": True, "strategy": "file_blocks", "modified_files": modified_files, "errors": [], } # --------------------------------------------------------------------------- # Public entry point: apply_ai_changes # --------------------------------------------------------------------------- def apply_ai_changes( project_root: str, generated_text: str, allowed_files: Optional[Sequence[str]] = None, ) -> Dict[str, Any]: """Apply AI-generated project modifications back onto the extracted ZIP. Priority (v4.5): 1. ```linepatch``` blocks — surgical, line-range-locked edits. 2. ```diff``` / ```patch``` blocks — unified diff via ``git apply``. 3. ``# FILE:`` fenced blocks — full-file rewrite (least trusted; should only be used for brand-new files). When ``allowed_files`` is provided, Diff Guard blocks any patch that tries to touch files outside the GraphMapper-derived edit scope. """ text = generated_text or "" if not text.strip(): return { "applied": False, "strategy": "none", "modified_files": [], "errors": ["empty generated output"], } guard = _enforce_diff_guard(text, allowed_files=allowed_files) if guard is not None: return guard # 1. Line-locked surgical patches — the preferred v4.5 path. linepatch_result = _apply_linepatch_blocks(project_root, text) if linepatch_result.get("applied"): return linepatch_result # 2. Unified diff / git patches. diff_result = _apply_diff_blocks(project_root, text) if diff_result.get("applied"): return diff_result # 3. Full-file rewrite blocks (compatibility fallback). file_block_result = _apply_file_blocks(project_root, text) if file_block_result.get("applied"): return file_block_result errors = [] errors.extend(linepatch_result.get("errors") or []) errors.extend(diff_result.get("errors") or []) errors.extend(file_block_result.get("errors") or []) return { "applied": False, "strategy": "none", "modified_files": [], "errors": errors or ["no applicable project modifications found"], } # --------------------------------------------------------------------------- # Utilities # --------------------------------------------------------------------------- def describe_bundle(bundle: Dict[str, Any], project_root: Optional[str] = None) -> Dict[str, Any]: """Return a lightweight, JSON-safe summary of an extracted ZIP bundle. Used by the Gradio UI (Code Graph tab) to show file counts, language mix, and top-level entries without dumping the whole preview blob a second time. """ root = project_root or bundle.get("project_root") or "" summary: Dict[str, Any] = { "zip_name": bundle.get("zip_name", ""), "archive_root_prefix": bundle.get("archive_root_prefix", ""), "project_root": root, "file_count": 0, "text_file_count": 0, "languages": {}, "top_level": [], } if not root or not os.path.isdir(root): return summary languages: Dict[str, int] = {} file_count = 0 text_count = 0 for rel_path, abs_path in _iter_project_files(root): file_count += 1 ext = abs_path.suffix.lower() or "(none)" languages[ext] = languages.get(ext, 0) + 1 if _looks_like_text(abs_path): text_count += 1 top_level: List[str] = [] try: for entry in sorted(os.listdir(root)): if entry in IGNORE_DIRS: continue top_level.append(entry + ("/" if os.path.isdir(os.path.join(root, entry)) else "")) except OSError: pass summary.update( { "file_count": file_count, "text_file_count": text_count, "languages": dict(sorted(languages.items(), key=lambda kv: -kv[1])[:12]), "top_level": top_level[:40], } ) return summary def cleanup_extraction(bundle_or_path: Any) -> bool: """Safely remove a previously extracted ZIP workspace. Accepts either the bundle dict returned by :func:`extract_uploaded_zip` or a raw ``extract_dir`` path string. Returns ``True`` when a directory was actually removed. Only touches directories under the OS temp dir or the DevAI workspace to avoid accidental recursive deletes. """ if not bundle_or_path: return False if isinstance(bundle_or_path, dict): path = bundle_or_path.get("extract_dir") or bundle_or_path.get("project_root") or "" else: path = str(bundle_or_path) if not path: return False path = os.path.abspath(path) safe_prefixes = [ os.path.abspath(tempfile.gettempdir()), os.path.abspath(os.path.join(tempfile.gettempdir(), "devai_zipintelligence")), os.path.abspath(os.environ.get("WORKSPACE_DIR", "/tmp/devai_workspace")), ] if not any(path.startswith(prefix + os.sep) or path == prefix for prefix in safe_prefixes): # Refuse to touch anything outside our known sandboxes. return False if not os.path.isdir(path): return False try: shutil.rmtree(path, ignore_errors=True) return not os.path.isdir(path) except Exception: return False def recompress_project_zip( project_root: str, output_dir: str, original_zip_name: Optional[str] = None, suffix: str = "zip_intelligence", archive_root_prefix: str = "", ) -> str: os.makedirs(output_dir, exist_ok=True) base_name = os.path.splitext(original_zip_name or "project.zip")[0] out_path = os.path.join(output_dir, f"{base_name}_{suffix}.zip") root = Path(project_root).resolve() with zipfile.ZipFile(out_path, "w", zipfile.ZIP_DEFLATED) as zf: for rel_path, abs_path in _iter_project_files(str(root)): if archive_root_prefix: arcname = f"{archive_root_prefix.rstrip('/')}/{rel_path}" else: arcname = rel_path zf.write(abs_path, arcname) return out_path # --------------------------------------------------------------------------- # v4.6 — Atomic Multi-file Patcher # --------------------------------------------------------------------------- # # Contract # -------- # If a task requires changes across N files, we MUST either: # # * apply every patch successfully → commit (all-or-nothing SUCCESS), OR # * revert the entire project to its pre-patch state on ANY failure → full # rollback (all-or-nothing FAILURE). # # The atomic wrapper works by snapshotting every file that appears in the # generated patch (plus their parent directory state so newly-created files # can be cleanly removed) BEFORE dispatching to the standard strategies # implemented in :func:`apply_ai_changes`. On failure, the snapshot is # restored byte-for-byte and any brand-new files are deleted. # # The snapshot is stored on disk under a per-transaction ``.devai_snapshot`` # directory inside ``project_root`` so we survive re-entrant calls. The # directory is always removed at the end (success or failure). class AtomicPatchError(RuntimeError): """Raised when an atomic multi-file patch could not be fully applied. The exception is intentionally lightweight: the recommended failure channel is the ``result`` dict returned by :func:`apply_ai_changes_atomic`, which already contains ``rolled_back`` and ``errors`` fields. """ def _collect_touched_files(text: str) -> List[str]: """Return every relative path referenced by ANY patch strategy. Used as the snapshot scope for atomic transactions. """ return _collect_candidate_modified_files(text) def _snapshot_files( project_root: str, files: Sequence[str], ) -> Dict[str, Any]: """Take a byte-level snapshot of *files* under *project_root*. Returns a manifest describing which files existed and where their pre-patch contents were stored on disk. The manifest is consumed by :func:`_rollback_from_snapshot`. """ root = Path(project_root).resolve() snap_dir = Path(tempfile.mkdtemp(prefix="devai_atomic_", dir=str(root))) manifest: Dict[str, Any] = { "snapshot_dir": str(snap_dir), "root": str(root), "entries": [], # list of {"rel": str, "existed": bool, "backup": str|None} } seen: set = set() for raw in files: try: rel = _safe_rel_path(raw) except ValueError: continue if rel in seen: continue seen.add(rel) abs_path = (root / rel).resolve() if root not in abs_path.parents and abs_path != root: # Refuse to snapshot anything outside the project. continue entry: Dict[str, Any] = {"rel": rel, "existed": abs_path.is_file(), "backup": None} if entry["existed"]: backup_path = snap_dir / rel backup_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(abs_path, backup_path) entry["backup"] = str(backup_path) manifest["entries"].append(entry) return manifest def _rollback_from_snapshot(manifest: Dict[str, Any]) -> List[str]: """Restore every file in *manifest* to its pre-patch state. Returns the list of relative paths that were rolled back (either restored or deleted). """ root = Path(manifest.get("root") or "").resolve() entries = manifest.get("entries") or [] rolled_back: List[str] = [] for entry in entries: rel = entry.get("rel") if not rel: continue abs_path = (root / rel).resolve() if root not in abs_path.parents and abs_path != root: continue if entry.get("existed"): backup = entry.get("backup") if backup and os.path.isfile(backup): try: abs_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(backup, abs_path) rolled_back.append(rel) except Exception: # Best-effort rollback — surface the failure via the # caller's ``errors`` channel rather than raising here. pass else: # The file was CREATED by this transaction — delete it if it # now exists on disk. if abs_path.exists(): try: abs_path.unlink() rolled_back.append(rel) # Prune empty parent directories that we implicitly # created (do NOT touch pre-existing folders). parent = abs_path.parent while parent != root and parent.is_dir() and not any(parent.iterdir()): parent.rmdir() parent = parent.parent except Exception: pass return rolled_back def _discard_snapshot(manifest: Dict[str, Any]) -> None: snap_dir = manifest.get("snapshot_dir") if isinstance(manifest, dict) else None if snap_dir and os.path.isdir(snap_dir): shutil.rmtree(snap_dir, ignore_errors=True) def _verify_all_touched_files_landed( project_root: str, result: Dict[str, Any], expected: Sequence[str], ) -> Optional[str]: """Confirm every file the AI *claimed* to modify actually exists. A patch may report ``applied=True`` yet skip individual files (e.g. when a linepatch range is out-of-bounds). For atomic guarantees we require every declared file to end up on disk after the transaction. """ reported = set(result.get("modified_files") or []) missing: List[str] = [] for rel in expected: try: clean = _safe_rel_path(rel) except ValueError: missing.append(rel) continue if clean not in reported: missing.append(clean) continue abs_path = Path(project_root) / clean if not abs_path.is_file(): missing.append(clean) if missing: return ( "atomic guard: the following files were declared in the patch " "but did not land on disk — " + ", ".join(sorted(set(missing))) ) return None def apply_ai_changes_atomic( project_root: str, generated_text: str, allowed_files: Optional[Sequence[str]] = None, ) -> Dict[str, Any]: """Atomic multi-file variant of :func:`apply_ai_changes`. Behaviour --------- 1. Enumerate every file referenced by the generated patch (across all three strategies — linepatch, diff, and ``# FILE:`` blocks). 2. Snapshot each of those files (and record whether they existed). 3. Delegate to :func:`apply_ai_changes`. 4. If ``applied`` is False, or any declared file did not land on disk, or an unexpected exception occurred, restore the snapshot and mark ``rolled_back=True``. 5. Otherwise, discard the snapshot and return the successful result augmented with ``atomic=True`` and ``rolled_back=False``. Returned dict extends the standard :func:`apply_ai_changes` shape with: ``atomic``: ``True`` — this response came from the atomic path. ``rolled_back``: whether a rollback was executed. ``touched_files``: files referenced by the patch (snapshot scope). ``rolled_back_files``: files whose state was actually reverted. """ text = generated_text or "" if not text.strip(): return { "applied": False, "strategy": "none", "modified_files": [], "errors": ["empty generated output"], "atomic": True, "rolled_back": False, "touched_files": [], "rolled_back_files": [], } # Diff Guard runs first — no snapshot needed if the patch is out of scope. guard = _enforce_diff_guard(text, allowed_files=allowed_files) if guard is not None: guard.update({ "atomic": True, "rolled_back": False, "touched_files": _collect_touched_files(text), "rolled_back_files": [], }) return guard touched = _collect_touched_files(text) manifest = _snapshot_files(project_root, touched) try: result = apply_ai_changes(project_root, text, allowed_files=allowed_files) except Exception as exc: # pragma: no cover — defensive rolled = _rollback_from_snapshot(manifest) _discard_snapshot(manifest) return { "applied": False, "strategy": "atomic_rollback", "modified_files": [], "errors": [f"atomic dispatch raised: {exc}"], "atomic": True, "rolled_back": True, "touched_files": touched, "rolled_back_files": rolled, } if not result.get("applied"): rolled = _rollback_from_snapshot(manifest) _discard_snapshot(manifest) result_out = dict(result) result_out.update({ "atomic": True, "rolled_back": True, "touched_files": touched, "rolled_back_files": rolled, }) return result_out # Multi-file completeness check: every declared file must have landed. completeness_error = _verify_all_touched_files_landed( project_root, result, touched ) if completeness_error is not None: rolled = _rollback_from_snapshot(manifest) _discard_snapshot(manifest) result_out = dict(result) errors = list(result_out.get("errors") or []) errors.append(completeness_error) result_out.update({ "applied": False, "errors": errors, "atomic": True, "rolled_back": True, "touched_files": touched, "rolled_back_files": rolled, }) return result_out # Success — discard the snapshot. _discard_snapshot(manifest) result_out = dict(result) result_out.update({ "atomic": True, "rolled_back": False, "touched_files": touched, "rolled_back_files": [], }) return result_out __all__ = [ "extract_uploaded_zip", "read_project_file", "list_project_files", "build_file_tree", "build_zip_ai_context", "preview_project_contents", "describe_bundle", "cleanup_extraction", "apply_ai_changes", "apply_ai_changes_atomic", "AtomicPatchError", "recompress_project_zip", ]