Spaces:
Sleeping
Sleeping
| """ | |
| 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: <start>-<end>`` 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<lang>[\w.+-]*)\n(?P<body>.*?)```", 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 | |
| # --- | |
| # <verbatim replacement content for lines 42..57 inclusive> | |
| # ``` | |
| # | |
| # 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<path>[^\n\r]+)\s*$", re.I | re.M) | |
| _LP_LINES_RE = re.compile( | |
| r"^\s*LINES:\s*(?P<start>\d+)\s*-\s*(?P<end>\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", | |
| ] | |