Improve tool output quality: ANSI stripping, tail-biased truncation, noise suppression

- Strip ANSI escape sequences from bash/job output (sandbox, local, jobs)
- Switch to tail-biased truncation (25/75 head/tail split) so errors and
final results are preserved instead of setup noise
- Spill full output to temp file when truncated, with metadata telling
the agent how to read specific sections via offset/limit
- Add UV_NO_PROGRESS=1 and PYTHONWARNINGS env vars to suppress install
progress and deprecation warnings at source
- Add atomic file writes, fuzzy edit matching, and Python validation
to sandbox server and local tools

agent/tools/edit_utils.py

@@ -0,0 +1,280 @@
+"""
+Shared utilities for file editing tools — fuzzy matching, syntax validation,
+and richer edit operations.
+
+Used by both local_tools.py and the embedded sandbox server.
+"""
+
+from __future__ import annotations
+
+# ── Unicode normalization map ────────────────────────────────────────────
+
+UNICODE_MAP = {
+    "\u2013": "-",   # en-dash
+    "\u2014": "-",   # em-dash
+    "\u2212": "-",   # minus sign
+    "\u2018": "'",   # left single quote
+    "\u2019": "'",   # right single quote
+    "\u201c": '"',   # left double quote
+    "\u201d": '"',   # right double quote
+    "\u00a0": " ",   # non-breaking space
+    "\u2003": " ",   # em space
+    "\u2002": " ",   # en space
+    "\u200b": "",    # zero-width space
+    "\ufeff": "",    # BOM
+}
+
+
+def _normalize_unicode(s: str) -> str:
+    return "".join(UNICODE_MAP.get(c, c) for c in s)
+
+
+# ── 4-pass fuzzy matching ────────────────────────────────────────────────
+
+
+def fuzzy_find(content: str, pattern: str) -> tuple[int | None, str | None]:
+    """Find *pattern* in *content* with increasingly relaxed matching.
+
+    Returns (start_index_in_original_content, match_note) or (None, None).
+    The index always refers to the *original* content string so callers can
+    use ``content[idx : idx + len(matched_text)]`` for replacement.
+
+    Strategy (mirrors Codex):
+      1. Exact match
+      2. Right-trim each line (trailing whitespace)
+      3. Both-sides trim (all surrounding whitespace per line)
+      4. Unicode normalization on top of both-sides trim
+    """
+    # Pass 1 — exact
+    if pattern in content:
+        return content.index(pattern), None
+
+    # Helper: build a line-stripped version *and* a mapping from stripped
+    # positions back to original positions.  We need this so callers can
+    # apply the replacement on the original content, not the stripped copy.
+
+    def _build_stripped(text: str, strip_fn):
+        """Return (stripped_text, line_start_map).
+
+        line_start_map[i] = original byte offset of the start of line i.
+        """
+        orig_lines = text.split("\n")
+        stripped_lines = [strip_fn(l) for l in orig_lines]
+        return "\n".join(stripped_lines), orig_lines, stripped_lines
+
+    # Pass 2 — right-trim
+    c_rt, c_orig_lines, c_rt_lines = _build_stripped(content, str.rstrip)
+    p_rt = "\n".join(l.rstrip() for l in pattern.split("\n"))
+    idx = c_rt.find(p_rt)
+    if idx != -1:
+        orig_idx = _map_back(idx, c_orig_lines, c_rt_lines)
+        return orig_idx, "(matched after trimming trailing whitespace)"
+
+    # Pass 3 — both-sides trim
+    c_st, _, c_st_lines = _build_stripped(content, str.strip)
+    p_st = "\n".join(l.strip() for l in pattern.split("\n"))
+    idx = c_st.find(p_st)
+    if idx != -1:
+        orig_idx = _map_back(idx, c_orig_lines, c_st_lines)
+        return orig_idx, "(matched after trimming whitespace)"
+
+    # Pass 4 — unicode normalization + both-sides trim
+    c_norm = _normalize_unicode(c_st)
+    p_norm = _normalize_unicode(p_st)
+    idx = c_norm.find(p_norm)
+    if idx != -1:
+        orig_idx = _map_back(idx, c_orig_lines, c_st_lines)
+        return orig_idx, "(matched after unicode normalization)"
+
+    return None, None
+
+
+def _map_back(
+    stripped_idx: int,
+    orig_lines: list[str],
+    stripped_lines: list[str],
+) -> int:
+    """Map a character index in the stripped/joined text back to the original text."""
+    # Walk through stripped lines to find which line the index falls on
+    pos = 0
+    for i, sl in enumerate(stripped_lines):
+        line_end = pos + len(sl)
+        if stripped_idx <= line_end:
+            col_in_stripped = stripped_idx - pos
+            # Find where this stripped line's content starts in the original line
+            ol = orig_lines[i]
+            # The stripped line is a subset of the original line; find its offset
+            lstripped = len(ol) - len(ol.lstrip())
+            orig_col = lstripped + col_in_stripped
+            # Compute absolute position in original text
+            orig_pos = sum(len(orig_lines[j]) + 1 for j in range(i)) + orig_col
+            return orig_pos
+        pos = line_end + 1  # +1 for the \n
+    # Fallback: return 0 (shouldn't happen if idx is valid)
+    return 0
+
+
+def fuzzy_find_original_match(content: str, pattern: str) -> tuple[str | None, str | None]:
+    """Find the *original* text in content that matches pattern fuzzily.
+
+    Returns (original_matched_text, match_note) or (None, None).
+    This extracts the exact substring from the original content that
+    corresponds to the fuzzy match, preserving its original whitespace/unicode.
+    """
+    if pattern in content:
+        return pattern, None
+
+    idx, note = fuzzy_find(content, pattern)
+    if idx is None:
+        return None, None
+
+    # We need to find the original text span that corresponds to the match.
+    # The match covers len(pattern) worth of *logical* content.
+    # Count how many original lines the pattern spans.
+    pattern_lines = pattern.split("\n")
+    n_lines = len(pattern_lines)
+
+    # Find which original line the match starts on
+    orig_lines = content.split("\n")
+    char_pos = 0
+    start_line = 0
+    for i, ol in enumerate(orig_lines):
+        if char_pos + len(ol) >= idx:
+            start_line = i
+            break
+        char_pos += len(ol) + 1
+
+    end_line = min(start_line + n_lines, len(orig_lines))
+    # Extract the original lines that were matched
+    matched_lines = orig_lines[start_line:end_line]
+    original_text = "\n".join(matched_lines)
+    return original_text, note
+
+
+# ── Richer edit operations ───────────────────────────────────────────────
+
+
+def apply_edit(
+    content: str,
+    old_str: str,
+    new_str: str,
+    mode: str = "replace",
+    replace_all: bool = False,
+) -> tuple[str, int, str | None]:
+    """Apply an edit operation to content.
+
+    Modes:
+      - replace: replace first occurrence (or all if replace_all=True)
+      - replace_all: replace all occurrences (alias)
+      - append_after: insert new_str after old_str
+      - prepend_before: insert new_str before old_str
+
+    Returns (new_content, num_replacements, fuzzy_note).
+    Raises ValueError if old_str not found.
+    """
+    if mode == "replace_all":
+        replace_all = True
+        mode = "replace"
+
+    # Try exact match first, then fuzzy
+    fuzzy_note = None
+    if old_str not in content:
+        original_match, fuzzy_note = fuzzy_find_original_match(content, old_str)
+        if original_match is None:
+            raise ValueError("old_str not found in file.")
+        old_str = original_match
+
+    count = content.count(old_str)
+
+    if mode == "replace":
+        if count > 1 and not replace_all:
+            raise ValueError(
+                f"old_str appears {count} times. Use replace_all=true to replace all, "
+                "or provide a more specific old_str."
+            )
+        if replace_all:
+            new_content = content.replace(old_str, new_str)
+            return new_content, count, fuzzy_note
+        else:
+            new_content = content.replace(old_str, new_str, 1)
+            return new_content, 1, fuzzy_note
+
+    elif mode == "append_after":
+        if replace_all:
+            new_content = content.replace(old_str, old_str + new_str)
+            return new_content, count, fuzzy_note
+        else:
+            idx = content.index(old_str) + len(old_str)
+            new_content = content[:idx] + new_str + content[idx:]
+            return new_content, 1, fuzzy_note
+
+    elif mode == "prepend_before":
+        if replace_all:
+            new_content = content.replace(old_str, new_str + old_str)
+            return new_content, count, fuzzy_note
+        else:
+            idx = content.index(old_str)
+            new_content = content[:idx] + new_str + content[idx:]
+            return new_content, 1, fuzzy_note
+
+    else:
+        raise ValueError(f"Unknown edit mode: {mode}. Use replace, append_after, or prepend_before.")
+
+
+# ── Syntax validation (Python) ───────────────────────────────────────────
+
+
+def validate_python(content: str, path: str = "") -> list[str]:
+    """Lightweight post-write validation for Python files.
+
+    Returns a list of warning strings (empty = all good).
+    Never raises — validation failures are advisory only.
+    """
+    import ast
+    import importlib
+
+    warnings = []
+
+    # 1. Syntax check via ast.parse
+    try:
+        tree = ast.parse(content)
+    except SyntaxError as e:
+        warnings.append(f"Python syntax error at line {e.lineno}: {e.msg}")
+        return warnings  # can't do import checks on broken syntax
+
+    # 2. Validate imports resolve
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ImportFrom):
+            if node.module:
+                try:
+                    mod = importlib.import_module(node.module)
+                    for alias in node.names:
+                        if alias.name != "*" and not hasattr(mod, alias.name):
+                            warnings.append(
+                                f"Import warning: '{alias.name}' not found in '{node.module}' (line {node.lineno})"
+                            )
+                except ImportError as e:
+                    warnings.append(f"Import error: {e} (line {node.lineno})")
+                except Exception:
+                    pass  # skip non-importable modules (e.g. project-local)
+        elif isinstance(node, ast.Import):
+            for alias in node.names:
+                try:
+                    importlib.import_module(alias.name)
+                except ImportError as e:
+                    warnings.append(f"Import error: {e} (line {node.lineno})")
+                except Exception:
+                    pass
+
+    # 3. Training script heuristics
+    if any(kw in content for kw in ("TrainingArguments", "SFTConfig", "DPOConfig", "GRPOConfig")):
+        if "push_to_hub" not in content:
+            warnings.append(
+                "Training script warning: no 'push_to_hub' found — model may be lost when job ends"
+            )
+        if "hub_model_id" not in content:
+            warnings.append(
+                "Training script warning: no 'hub_model_id' found"
+            )
+
+    return warnings

agent/tools/jobs_tool.py

@@ -117,11 +117,19 @@ def _filter_uv_install_output(logs: list[str]) -> list[str]:
     return logs
 
 
+_ANSI_RE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07')
+
+
+def _strip_ansi(text: str) -> str:
+    return _ANSI_RE.sub('', text)
+
+
 _DEFAULT_ENV = {
     "HF_HUB_DISABLE_PROGRESS_BARS": "1",
     "TQDM_DISABLE": "1",
     "TRANSFORMERS_VERBOSITY": "warning",
     "HF_HUB_ENABLE_HF_TRANSFER": "1",
+    "UV_NO_PROGRESS": "1",
 }
 
 
@@ -580,7 +588,7 @@ class HfJobsTool:
             filtered_logs = _filter_uv_install_output(all_logs)
 
             # Format all logs for the agent
-            log_text = "\n".join(filtered_logs) if filtered_logs else "(no logs)"
+            log_text = _strip_ansi("\n".join(filtered_logs)) if filtered_logs else "(no logs)"
 
             response = f"""{job_type} job completed!
 
@@ -657,7 +665,7 @@ class HfJobsTool:
                     "resultsShared": 0,
                 }
 
-            log_text = "\n".join(logs)
+            log_text = _strip_ansi("\n".join(logs))
             return {
                 "formatted": f"**Logs for {job_id}:**\n\n```\n{log_text}\n```",
                 "totalResults": 1,

agent/tools/local_tools.py

@@ -8,18 +8,78 @@ subprocess/pathlib instead of going through a remote sandbox.
 
 from __future__ import annotations
 
+import os
+import re
 import subprocess
+import tempfile
 from pathlib import Path
 from typing import Any
 
 from agent.tools.sandbox_client import Sandbox
 
-MAX_OUTPUT_CHARS = 30_000
+MAX_OUTPUT_CHARS = 25_000
 MAX_LINE_LENGTH = 2000
 DEFAULT_READ_LINES = 2000
 DEFAULT_TIMEOUT = 120
 MAX_TIMEOUT = 600
 
+_ANSI_RE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07')
+
+
+def _atomic_write(path: Path, content: str) -> None:
+    """Write file atomically via temp file + os.replace().
+
+    Ensures the file is never left in a partial/corrupted state — it's either
+    the old content or the new content, never half-written.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd = None
+    tmp_path = None
+    try:
+        fd, tmp_path = tempfile.mkstemp(dir=path.parent, suffix=".tmp")
+        os.write(fd, content.encode("utf-8"))
+        os.fsync(fd)
+        os.close(fd)
+        fd = None
+        os.replace(tmp_path, str(path))
+        tmp_path = None  # successfully replaced, nothing to clean up
+    finally:
+        if fd is not None:
+            os.close(fd)
+        if tmp_path is not None:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+
+
+def _strip_ansi(text: str) -> str:
+    return _ANSI_RE.sub('', text)
+
+
+def _truncate_output(output: str, max_chars: int = MAX_OUTPUT_CHARS, head_ratio: float = 0.25) -> str:
+    """Tail-biased truncation with temp file spillover for full output access."""
+    if len(output) <= max_chars:
+        return output
+    # Write full output to temp file so LLM can read specific sections
+    spill_path = None
+    try:
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', prefix='bash_output_', delete=False) as f:
+            f.write(output)
+            spill_path = f.name
+    except Exception:
+        pass
+    head_budget = int(max_chars * head_ratio)
+    tail_budget = max_chars - head_budget
+    head = output[:head_budget]
+    tail = output[-tail_budget:]
+    total = len(output)
+    omitted = total - max_chars
+    meta = f"\n\n... ({omitted:,} of {total:,} chars omitted, showing first {head_budget:,} + last {tail_budget:,}) ...\n"
+    if spill_path:
+        meta += f"Full output saved to {spill_path} — use the read tool with offset/limit to inspect specific sections.\n"
+    return head + meta + tail
+
 
 # ── Handlers ────────────────────────────────────────────────────────────
 
@@ -38,9 +98,8 @@ async def _bash_handler(args: dict[str, Any], **_kw) -> tuple[str, bool]:
             cwd=work_dir,
             timeout=timeout,
         )
-        output = result.stdout + result.stderr
-        if len(output) > MAX_OUTPUT_CHARS:
-            output = output[:MAX_OUTPUT_CHARS] + "\n... (output truncated)"
+        output = _strip_ansi(result.stdout + result.stderr)
+        output = _truncate_output(output)
         if not output.strip():
             output = "(no output)"
         return output, result.returncode == 0
@@ -83,18 +142,27 @@ async def _write_handler(args: dict[str, Any], **_kw) -> tuple[str, bool]:
         return "No path provided.", False
     p = Path(file_path)
     try:
-        p.parent.mkdir(parents=True, exist_ok=True)
-        p.write_text(content)
-        return f"Wrote {len(content)} bytes to {file_path}", True
+        _atomic_write(p, content)
+        msg = f"Wrote {len(content)} bytes to {file_path}"
+        # Syntax validation for Python files
+        if p.suffix == ".py":
+            from agent.tools.edit_utils import validate_python
+            warnings = validate_python(content, file_path)
+            if warnings:
+                msg += "\n\nValidation warnings:\n" + "\n".join(f"  ⚠ {w}" for w in warnings)
+        return msg, True
     except Exception as e:
         return f"write error: {e}", False
 
 
 async def _edit_handler(args: dict[str, Any], **_kw) -> tuple[str, bool]:
+    from agent.tools.edit_utils import apply_edit, validate_python
+
     file_path = args.get("path", "")
     old_str = args.get("old_str", "")
     new_str = args.get("new_str", "")
     replace_all = args.get("replace_all", False)
+    mode = args.get("mode", "replace")
 
     if not file_path:
         return "No path provided.", False
@@ -110,23 +178,27 @@ async def _edit_handler(args: dict[str, Any], **_kw) -> tuple[str, bool]:
     except Exception as e:
         return f"edit read error: {e}", False
 
-    count = text.count(old_str)
-    if count == 0:
-        return "old_str not found in file.", False
-    if count > 1 and not replace_all:
-        return (
-            f"old_str appears {count} times. Use replace_all=true to replace all, "
-            "or provide a more specific old_str."
-        ), False
+    try:
+        new_text, replacements, fuzzy_note = apply_edit(
+            text, old_str, new_str, mode=mode, replace_all=replace_all
+        )
+    except ValueError as e:
+        return str(e), False
 
-    new_text = text.replace(old_str, new_str) if replace_all else text.replace(old_str, new_str, 1)
     try:
-        p.write_text(new_text)
+        _atomic_write(p, new_text)
     except Exception as e:
         return f"edit write error: {e}", False
 
-    replacements = count if replace_all else 1
-    return f"Edited {file_path} ({replacements} replacement{'s' if replacements > 1 else ''})", True
+    msg = f"Edited {file_path} ({replacements} replacement{'s' if replacements > 1 else ''})"
+    if fuzzy_note:
+        msg += f" {fuzzy_note}"
+    # Syntax validation for Python files
+    if p.suffix == ".py":
+        warnings = validate_python(new_text, file_path)
+        if warnings:
+            msg += "\n\nValidation warnings:\n" + "\n".join(f"  ⚠ {w}" for w in warnings)
+    return msg, True
 
 
 # ── Public API ──────────────────────────────────────────────────────────

agent/tools/sandbox_client.py

@@ -56,7 +56,7 @@ HARDWARE_OPTIONS = [
     "a10g-large",
     "a100-large",
 ]
-OUTPUT_LIMIT = 30000
+OUTPUT_LIMIT = 25000
 LINE_LIMIT = 2000
 DEFAULT_READ_LIMIT = 2000
 DEFAULT_TIMEOUT = 240
@@ -85,7 +85,9 @@ ENV HOME=/home/user \\
     PIP_USER=1 \\
     HF_HUB_DISABLE_PROGRESS_BARS=1 \\
     TQDM_DISABLE=1 \\
-    HF_HUB_ENABLE_HF_TRANSFER=1
+    HF_HUB_ENABLE_HF_TRANSFER=1 \\
+    UV_NO_PROGRESS=1 \\
+    PYTHONWARNINGS=ignore::DeprecationWarning
 
 WORKDIR /app
 COPY --chown=user . /app
@@ -97,12 +99,61 @@ CMD ["python", "sandbox_server.py"]
 
 _SANDBOX_SERVER = '''\
 """Minimal FastAPI server for sandbox operations."""
-import os, subprocess, pathlib, signal, threading
+import os, subprocess, pathlib, signal, threading, re, tempfile
 from fastapi import FastAPI
 from pydantic import BaseModel
 from typing import Optional
 import uvicorn
 
+_ANSI_RE = re.compile(r'\\x1b\\[[0-9;]*[a-zA-Z]|\\x1b\\].*?\\x07')
+
+def _strip_ansi(text: str) -> str:
+    return _ANSI_RE.sub('', text)
+
+def _truncate_output(output: str, max_chars: int = 25000, head_ratio: float = 0.25) -> str:
+    if len(output) <= max_chars:
+        return output
+    # Write full output to temp file so LLM can read specific sections
+    spill_path = None
+    try:
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', prefix='bash_output_', dir='/tmp', delete=False) as f:
+            f.write(output)
+            spill_path = f.name
+    except Exception:
+        pass
+    head_budget = int(max_chars * head_ratio)
+    tail_budget = max_chars - head_budget
+    head = output[:head_budget]
+    tail = output[-tail_budget:]
+    total = len(output)
+    omitted = total - max_chars
+    meta = f"\\n\\n... ({omitted:,} of {total:,} chars omitted, showing first {head_budget:,} + last {tail_budget:,}) ...\\n"
+    if spill_path:
+        meta += f"Full output saved to {spill_path} — use the read tool with offset/limit to inspect specific sections.\\n"
+    return head + meta + tail
+
+def _atomic_write(path: pathlib.Path, content: str):
+    """Write atomically: temp file + fsync + os.replace."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd = None
+    tmp_path = None
+    try:
+        fd, tmp_path = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp")
+        os.write(fd, content.encode("utf-8"))
+        os.fsync(fd)
+        os.close(fd)
+        fd = None
+        os.replace(tmp_path, str(path))
+        tmp_path = None
+    finally:
+        if fd is not None:
+            os.close(fd)
+        if tmp_path is not None:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+
 app = FastAPI()
 
 # Track active bash processes so they can be killed on cancel
@@ -128,10 +179,123 @@ class EditReq(BaseModel):
     old_str: str
     new_str: str
     replace_all: bool = False
+    mode: str = "replace"
 
 class ExistsReq(BaseModel):
     path: str
 
+# ── Fuzzy matching & edit utilities (embedded) ──
+
+UNICODE_MAP = {
+    "\\u2013": "-", "\\u2014": "-", "\\u2212": "-",
+    "\\u2018": "'", "\\u2019": "'",
+    "\\u201c": \'"\', "\\u201d": \'"\',
+    "\\u00a0": " ", "\\u2003": " ", "\\u2002": " ",
+    "\\u200b": "", "\\ufeff": "",
+}
+
+def _normalize_unicode(s):
+    return "".join(UNICODE_MAP.get(c, c) for c in s)
+
+def _fuzzy_find_original(content, pattern):
+    """Find the original text in content that matches pattern fuzzily."""
+    if pattern in content:
+        return pattern, None
+    # Pass 2: right-trim
+    c_lines = content.split("\\n")
+    c_rt = "\\n".join(l.rstrip() for l in c_lines)
+    p_rt = "\\n".join(l.rstrip() for l in pattern.split("\\n"))
+    if p_rt in c_rt:
+        idx = c_rt.index(p_rt)
+        start_line = c_rt[:idx].count("\\n")
+        n_lines = p_rt.count("\\n") + 1
+        matched = "\\n".join(c_lines[start_line:start_line + n_lines])
+        return matched, "(matched after trimming trailing whitespace)"
+    # Pass 3: both-sides trim
+    c_st = "\\n".join(l.strip() for l in c_lines)
+    p_st = "\\n".join(l.strip() for l in pattern.split("\\n"))
+    if p_st in c_st:
+        idx = c_st.index(p_st)
+        start_line = c_st[:idx].count("\\n")
+        n_lines = p_st.count("\\n") + 1
+        matched = "\\n".join(c_lines[start_line:start_line + n_lines])
+        return matched, "(matched after trimming whitespace)"
+    # Pass 4: unicode normalization
+    c_norm = _normalize_unicode(c_st)
+    p_norm = _normalize_unicode(p_st)
+    if p_norm in c_norm:
+        idx = c_norm.index(p_norm)
+        start_line = c_norm[:idx].count("\\n")
+        n_lines = p_norm.count("\\n") + 1
+        matched = "\\n".join(c_lines[start_line:start_line + n_lines])
+        return matched, "(matched after unicode normalization)"
+    return None, None
+
+def _apply_edit(content, old_str, new_str, mode="replace", replace_all=False):
+    """Apply edit. Returns (new_content, count, fuzzy_note) or raises ValueError."""
+    if mode == "replace_all":
+        replace_all = True
+        mode = "replace"
+    fuzzy_note = None
+    if old_str not in content:
+        matched, fuzzy_note = _fuzzy_find_original(content, old_str)
+        if matched is None:
+            raise ValueError("old_str not found in file.")
+        old_str = matched
+    count = content.count(old_str)
+    if mode == "replace":
+        if count > 1 and not replace_all:
+            raise ValueError(f"old_str appears {count} times. Use replace_all=true or provide more context.")
+        if replace_all:
+            return content.replace(old_str, new_str), count, fuzzy_note
+        return content.replace(old_str, new_str, 1), 1, fuzzy_note
+    elif mode == "append_after":
+        if replace_all:
+            return content.replace(old_str, old_str + new_str), count, fuzzy_note
+        idx = content.index(old_str) + len(old_str)
+        return content[:idx] + new_str + content[idx:], 1, fuzzy_note
+    elif mode == "prepend_before":
+        if replace_all:
+            return content.replace(old_str, new_str + old_str), count, fuzzy_note
+        idx = content.index(old_str)
+        return content[:idx] + new_str + content[idx:], 1, fuzzy_note
+    raise ValueError(f"Unknown mode: {mode}")
+
+def _validate_python(content, path=""):
+    """Lightweight Python validation. Returns list of warning strings."""
+    import ast as _ast, importlib as _il
+    warnings = []
+    try:
+        tree = _ast.parse(content)
+    except SyntaxError as e:
+        warnings.append(f"Python syntax error at line {e.lineno}: {e.msg}")
+        return warnings
+    for node in _ast.walk(tree):
+        if isinstance(node, _ast.ImportFrom) and node.module:
+            try:
+                mod = _il.import_module(node.module)
+                for alias in node.names:
+                    if alias.name != "*" and not hasattr(mod, alias.name):
+                        warnings.append(f"Import warning: '{alias.name}' not found in '{node.module}' (line {node.lineno})")
+            except ImportError as e:
+                warnings.append(f"Import error: {e} (line {node.lineno})")
+            except Exception:
+                pass
+        elif isinstance(node, _ast.Import):
+            for alias in node.names:
+                try:
+                    _il.import_module(alias.name)
+                except ImportError as e:
+                    warnings.append(f"Import error: {e} (line {node.lineno})")
+                except Exception:
+                    pass
+    if any(kw in content for kw in ("TrainingArguments", "SFTConfig", "DPOConfig", "GRPOConfig")):
+        if "push_to_hub" not in content:
+            warnings.append("Training script warning: no \'push_to_hub\' found")
+        if "hub_model_id" not in content:
+            warnings.append("Training script warning: no \'hub_model_id\' found")
+    return warnings
+
 @app.get("/api/health")
 def health():
     return {"status": "ok"}
@@ -147,9 +311,8 @@ def bash(req: BashReq):
             _active_procs[proc.pid] = proc
         try:
             stdout, stderr = proc.communicate(timeout=req.timeout)
-            output = stdout + stderr
-            if len(output) > 30000:
-                output = output[:30000] + "\\n... (truncated)"
+            output = _strip_ansi(stdout + stderr)
+            output = _truncate_output(output)
             return {"success": proc.returncode == 0, "output": output, "error": "" if proc.returncode == 0 else f"Exit code {proc.returncode}"}
         except subprocess.TimeoutExpired:
             try:
@@ -203,9 +366,13 @@ def read(req: ReadReq):
 def write(req: WriteReq):
     try:
         p = pathlib.Path(req.path)
-        p.parent.mkdir(parents=True, exist_ok=True)
-        p.write_text(req.content)
-        return {"success": True, "output": f"Wrote {len(req.content)} bytes to {req.path}", "error": ""}
+        _atomic_write(p, req.content)
+        msg = f"Wrote {len(req.content)} bytes to {req.path}"
+        if p.suffix == ".py":
+            warnings = _validate_python(req.content, req.path)
+            if warnings:
+                msg += "\\n\\nValidation warnings:\\n" + "\\n".join(f"  ! {w}" for w in warnings)
+        return {"success": True, "output": msg, "error": ""}
     except Exception as e:
         return {"success": False, "output": "", "error": str(e)}
 
@@ -216,16 +383,23 @@ def edit(req: EditReq):
         if not p.exists():
             return {"success": False, "output": "", "error": f"File not found: {req.path}"}
         content = p.read_text()
-        if req.old_str not in content:
-            return {"success": False, "output": "", "error": f"old_str not found in {req.path}"}
-        if not req.replace_all and content.count(req.old_str) > 1:
-            return {"success": False, "output": "", "error": f"old_str appears {content.count(req.old_str)} times. Use replace_all=true or provide more context."}
-        if req.replace_all:
-            new_content = content.replace(req.old_str, req.new_str)
-        else:
-            new_content = content.replace(req.old_str, req.new_str, 1)
-        p.write_text(new_content)
-        return {"success": True, "output": f"Edited {req.path}", "error": ""}
+        if req.old_str == req.new_str:
+            return {"success": False, "output": "", "error": "old_str and new_str must differ."}
+        try:
+            new_content, count, fuzzy_note = _apply_edit(
+                content, req.old_str, req.new_str, mode=req.mode, replace_all=req.replace_all
+            )
+        except ValueError as e:
+            return {"success": False, "output": "", "error": str(e)}
+        _atomic_write(p, new_content)
+        msg = f"Edited {req.path} ({count} replacement{'s' if count > 1 else ''})"
+        if fuzzy_note:
+            msg += f" {fuzzy_note}"
+        if p.suffix == ".py":
+            warnings = _validate_python(new_content, req.path)
+            if warnings:
+                msg += "\\n\\nValidation warnings:\\n" + "\\n".join(f"  ! {w}" for w in warnings)
+        return {"success": True, "output": msg, "error": ""}
     except Exception as e:
         return {"success": False, "output": "", "error": str(e)}
 
@@ -605,7 +779,8 @@ class Sandbox:
         return result
 
     def edit(
-        self, path: str, old_str: str, new_str: str, *, replace_all: bool = False
+        self, path: str, old_str: str, new_str: str, *, replace_all: bool = False,
+        mode: str = "replace",
     ) -> ToolResult:
         if old_str == new_str:
             return ToolResult(success=False, error="old_str and new_str are identical.")
@@ -621,6 +796,7 @@ class Sandbox:
                 "old_str": old_str,
                 "new_str": new_str,
                 "replace_all": replace_all,
+                "mode": mode,
             },
         )
 
@@ -731,7 +907,12 @@ class Sandbox:
         },
         "edit": {
             "description": (
-                "Targeted edit via exact string replacement.\n"
+                "Targeted edit via string replacement with fuzzy matching fallback.\n"
+                "\n"
+                "Modes:\n"
+                "- replace (default): replace first occurrence of old_str with new_str.\n"
+                "- append_after: insert new_str immediately after old_str (old_str is kept).\n"
+                "- prepend_before: insert new_str immediately before old_str (old_str is kept).\n"
                 "\n"
                 "Rules:\n"
                 "- old_str must appear EXACTLY once (unless replace_all is true).\n"
@@ -742,7 +923,9 @@ class Sandbox:
                 "- File MUST have been read this session (system enforced).\n"
                 "- Do NOT include line number prefixes in old_str/new_str.\n"
                 "\n"
-                "Use replace_all=true for batch operations like variable renaming."
+                "If exact match fails, the tool automatically tries trimmed/normalized matching.\n"
+                "Use replace_all=true for batch operations like variable renaming.\n"
+                "Use append_after/prepend_before to insert code without replacing existing code."
             ),
             "parameters": {
                 "type": "object",
@@ -755,14 +938,20 @@ class Sandbox:
                     },
                     "old_str": {
                         "type": "string",
-                        "description": "Exact text to find (must differ from new_str).",
+                        "description": "Text to find (fuzzy matching used as fallback).",
                     },
-                    "new_str": {"type": "string", "description": "Replacement text."},
+                    "new_str": {"type": "string", "description": "Replacement text (or text to insert for append_after/prepend_before)."},
                     "replace_all": {
                         "type": "boolean",
                         "description": "Replace all occurrences (default: false).",
                         "default": False,
                     },
+                    "mode": {
+                        "type": "string",
+                        "enum": ["replace", "append_after", "prepend_before"],
+                        "description": "Edit mode (default: replace).",
+                        "default": "replace",
+                    },
                 },
             },
         },
@@ -791,6 +980,7 @@ class Sandbox:
                 a["old_str"],
                 a["new_str"],
                 replace_all=a.get("replace_all", False),
+                mode=a.get("mode", "replace"),
             ),
         }
         fn = dispatch.get(name)