Update documentation: Add file diff rendering to REPL, bump version to 1.17.0, enhance AGENTS.md, and add test for resume interrupt

tai.py

@@ -3,6 +3,7 @@
 import argparse
 import base64
 import collections
+import difflib
 import getpass
 import glob
 import hashlib
@@ -36,7 +37,7 @@ try:
 except ImportError:
     readline = None
 
-VERSION = "1.16.0"
+VERSION = "1.17.0"
 WORKER_STEPS = 12
 CREATE_SKILL_STEPS = 40
 FORK_MAX_DEPTH = 2
@@ -436,6 +437,22 @@ def atomic_write(path, content):
     os.replace(tmp, path)
 
 
+def sealed_write(path, content):
+    atomic_write(path, content)
+    try:
+        os.chmod(path, 0o600)
+    except OSError:
+        pass
+
+
+def quarantine(path):
+    try:
+        if os.path.exists(path):
+            os.replace(path, "%s.corrupt-%d" % (path, int(time.time())))
+    except OSError:
+        pass
+
+
 def release_bump(script_file, part):
     with open(script_file, "r", encoding="utf-8") as handle:
         text = handle.read()
@@ -697,6 +714,145 @@ def run_live(argv, timeout=SHELL_TIMEOUT, shell=False, cwd=None, env=None, strea
     return proc.returncode, "".join(chunks), time.time() - started, timed_out
 
 
+DIFF_CONTEXT = 3
+DIFF_MAX_LINES = 120
+BG_ADD = "\033[48;5;28m"
+BG_DEL = "\033[48;5;88m"
+FG_ORANGE = "\033[38;5;214m"
+
+PY_KEYWORDS = ("False", "None", "True", "and", "as", "assert", "async", "await", "break", "class", "continue", "def", "del", "elif", "else", "except", "finally", "for", "from", "global", "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return", "try", "while", "with", "yield", "match", "case")
+
+PY_TOKEN_RE = re.compile(r"(?P<string>\"\"\"(?:\\.|[^\\])*?\"\"\"|'''(?:\\.|[^\\])*?'''|\"(?:\\.|[^\"\\\n])*\"|'(?:\\.|[^'\\\n])*')|(?P<comment>#[^\n]*)|(?P<number>\b\d[\d._]*(?:[eE][+-]?\d+)?[jJ]?\b)|(?P<keyword>\b(?:%s)\b)|(?P<decorator>@[A-Za-z_][\w.]*)|(?P<defname>(?<=def )[A-Za-z_]\w*|(?<=class )[A-Za-z_]\w*)" % "|".join(PY_KEYWORDS))
+
+
+def color_enabled():
+    return not os.environ.get("NO_COLOR") and sys.stdout.isatty()
+
+
+def highlight_python(line):
+    parts = []
+    pos = 0
+    for match in PY_TOKEN_RE.finditer(line):
+        parts.append(line[pos:match.start()])
+        kind = match.lastgroup
+        text = match.group(0)
+        if kind == "string":
+            parts.append(Ansi.YELLOW + text + Ansi.RESET)
+        elif kind == "comment":
+            parts.append(Ansi.GRAY + text + Ansi.RESET)
+        elif kind == "number":
+            parts.append(FG_ORANGE + text + Ansi.RESET)
+        elif kind == "keyword":
+            parts.append(Ansi.MAGENTA + text + Ansi.RESET)
+        else:
+            parts.append(Ansi.CYAN + text + Ansi.RESET)
+        pos = match.end()
+    parts.append(line[pos:])
+    return "".join(parts)
+
+
+def split_diff_lines(text):
+    if not text:
+        return []
+    lines = text.split("\n")
+    if lines and lines[-1] == "":
+        lines.pop()
+    return lines
+
+
+def render_diff(path, old, new, width=0, colors=True):
+    if new is not None and (old or "") == new:
+        return ""
+    old_lines = split_diff_lines(old or "")
+    new_lines = split_diff_lines(new) if new is not None else []
+    if not width:
+        try:
+            width = shutil.get_terminal_size().columns
+        except OSError:
+            width = 80
+    matcher = difflib.SequenceMatcher(None, old_lines, new_lines, autojunk=False)
+    rows = []
+    for group in matcher.get_grouped_opcodes(DIFF_CONTEXT):
+        if rows:
+            rows.append(("gap", 0, 0, ""))
+        for tag, i1, i2, j1, j2 in group:
+            if tag == "equal":
+                for offset in range(i2 - i1):
+                    rows.append(("ctx", i1 + offset + 1, j1 + offset + 1, old_lines[i1 + offset]))
+            elif tag == "delete":
+                for offset in range(i2 - i1):
+                    rows.append(("del", i1 + offset + 1, 0, old_lines[i1 + offset]))
+            elif tag == "insert":
+                for offset in range(j2 - j1):
+                    rows.append(("add", 0, j1 + offset + 1, new_lines[j1 + offset]))
+            else:
+                for offset in range(i2 - i1):
+                    rows.append(("del", i1 + offset + 1, 0, old_lines[i1 + offset]))
+                for offset in range(j2 - j1):
+                    rows.append(("add", 0, j1 + offset + 1, new_lines[j1 + offset]))
+    added = sum(1 for kind, _a, _b, _t in rows if kind == "add")
+    removed = sum(1 for kind, _a, _b, _t in rows if kind == "del")
+    if new is None:
+        head = "── %s (deleted, %d lines)" % (path, len(old_lines))
+    elif not old_lines:
+        head = "── %s (new file, %d lines)" % (path, len(new_lines))
+    else:
+        head = "── %s (+%d -%d)" % (path, added, removed)
+    python = str(path).lower().endswith(".py")
+    num_width = max(4, len(str(max(len(old_lines), len(new_lines), 1))))
+    code_width = max(20, width - num_width - 4)
+    fence = None
+    out = []
+    out.append(Ansi.DIM + "│ " + head + Ansi.RESET if colors else "│ " + head)
+    for kind, old_no, new_no, text in rows[:DIFF_MAX_LINES]:
+        if kind == "gap":
+            out.append(Ansi.DIM + "···" + Ansi.RESET if colors else "···")
+            continue
+        code = strip_ansi(text).expandtabs(8)
+        if fence is not None:
+            if python and colors:
+                code = Ansi.YELLOW + code + Ansi.RESET
+            if fence in text:
+                fence = None
+        elif python and colors:
+            for mark in ('"""', "'''"):
+                if text.count(mark) % 2 == 1:
+                    fence = mark
+                    code = Ansi.YELLOW + code + Ansi.RESET
+                    break
+            else:
+                code = highlight_python(code)
+        number = old_no if kind in ("ctx", "del") else new_no
+        gutter = str(number).rjust(num_width)
+        if not colors:
+            sign = {"ctx": " ", "del": "-", "add": "+"}[kind]
+            out.append("%s %s %s" % (gutter, sign, code[:code_width]))
+            continue
+        dimmed = Ansi.DIM + gutter + Ansi.RESET
+        if kind == "ctx":
+            out.append("%s   %s" % (dimmed, truncate_ansi(code, code_width)))
+        elif kind == "del":
+            row = "%s %s %s" % (dimmed, Ansi.RED + Ansi.BOLD + "-" + Ansi.RESET, truncate_ansi(code, code_width))
+            out.append(pad_tinted(style_span(row, BG_DEL), width))
+        else:
+            row = "%s %s %s" % (dimmed, Ansi.GREEN + Ansi.BOLD + "+" + Ansi.RESET, truncate_ansi(code, code_width))
+            out.append(pad_tinted(style_span(row, BG_ADD), width))
+    if len(rows) > DIFF_MAX_LINES:
+        note = "··· %d more lines hidden" % (len(rows) - DIFF_MAX_LINES)
+        out.append(Ansi.DIM + note + Ansi.RESET if colors else note)
+    return "\n".join(out)
+
+
+def pad_tinted(line, width):
+    plain = len(ANSI_RE.sub("", line))
+    if plain < width:
+        fill = " " * (width - plain)
+        if line.endswith(Ansi.RESET):
+            return line[:-len(Ansi.RESET)] + fill + Ansi.RESET
+        return line + fill
+    return line
+
+
 ANSI_RE = re.compile(r"\033\[[0-9;]*m")
 MD_BULLETS = ("•", "◦", "▪")
 
@@ -2100,17 +2256,23 @@ class Store:
 
     def load_system(self, name):
         path = self.profile_path(name)
-        if os.path.exists(path):
+        if not os.path.exists(path):
+            return None
+        try:
             with open(path, "r", encoding="utf-8") as handle:
                 return self.unseal_file(handle.read())
-        return None
+        except (OSError, ValueError):
+            quarantine(path)
+            return None
+        except SealError:
+            if not self.seal.enabled:
+                raise
+            quarantine(path)
+            return None
 
     def save_system(self, name, text):
-        path = self.profile_path(name)
-        with open(path, "w", encoding="utf-8") as handle:
-            handle.write(self.seal.lock(text))
         try:
-            os.chmod(path, 0o600)
+            sealed_write(self.profile_path(name), self.seal.lock(text))
         except OSError:
             pass
 
@@ -2124,14 +2286,46 @@ class Store:
             data = json.loads(self.unseal_file(raw))
             return [item for item in data if isinstance(item, dict) and item.get("role") in ("user", "assistant", "tool")]
         except (OSError, ValueError):
+            quarantine(path)
+            return []
+        except SealError:
+            if not self.seal.enabled:
+                raise
+            quarantine(path)
             return []
 
     def save_session(self, name, messages, bot="main"):
         path = self.session_path(name) if bot == "main" else self.bot_session_path(name, bot)
         try:
-            with open(path, "w", encoding="utf-8") as handle:
-                handle.write(self.seal.lock(json.dumps([item for item in messages if item.get("role") != "system"])))
-            os.chmod(path, 0o600)
+            sealed_write(path, self.seal.lock(json.dumps([item for item in messages if item.get("role") != "system"])))
+        except OSError:
+            pass
+
+    def turn_state_path(self, profile, bot="main"):
+        if bot == "main":
+            return os.path.join(self.config.profiles_dir, "%s.turn.json" % profile)
+        return os.path.join(self.config.profiles_dir, "%s.bot.%s.turn.json" % (profile, bot))
+
+    def load_turn_state(self, profile, bot="main"):
+        path = self.turn_state_path(profile, bot)
+        if not os.path.exists(path):
+            return {}
+        try:
+            with open(path, "r", encoding="utf-8") as handle:
+                data = json.loads(self.unseal_file(handle.read()))
+            return data if isinstance(data, dict) else {}
+        except (OSError, ValueError):
+            quarantine(path)
+            return {}
+        except SealError:
+            if not self.seal.enabled:
+                raise
+            quarantine(path)
+            return {}
+
+    def save_turn_state(self, profile, bot, state):
+        try:
+            sealed_write(self.turn_state_path(profile, bot), self.seal.lock(json.dumps(state)))
         except OSError:
             pass
 
@@ -2154,13 +2348,17 @@ class Store:
             bots = data.get("bots") if isinstance(data, dict) else None
             return bots if isinstance(bots, dict) else {}
         except (OSError, ValueError):
+            quarantine(path)
+            return {}
+        except SealError:
+            if not self.seal.enabled:
+                raise
+            quarantine(path)
             return {}
 
     def save_bots(self, profile, bots):
         try:
-            with open(self.bots_registry_path(profile), "w", encoding="utf-8") as handle:
-                handle.write(self.seal.lock(json.dumps({"bots": bots})))
-            os.chmod(self.bots_registry_path(profile), 0o600)
+            sealed_write(self.bots_registry_path(profile), self.seal.lock(json.dumps({"bots": bots})))
         except OSError:
             pass
 
@@ -2187,17 +2385,23 @@ class Store:
         if bot == "main":
             return self.load_system(profile)
         path = self.bot_system_path(profile, bot)
-        if os.path.exists(path):
+        if not os.path.exists(path):
+            return None
+        try:
             with open(path, "r", encoding="utf-8") as handle:
                 return self.unseal_file(handle.read())
-        return None
+        except (OSError, ValueError):
+            quarantine(path)
+            return None
+        except SealError:
+            if not self.seal.enabled:
+                raise
+            quarantine(path)
+            return None
 
     def save_bot_system(self, profile, bot, text):
-        path = self.bot_system_path(profile, bot)
-        with open(path, "w", encoding="utf-8") as handle:
-            handle.write(self.seal.lock(text))
         try:
-            os.chmod(path, 0o600)
+            sealed_write(self.bot_system_path(profile, bot), self.seal.lock(text))
         except OSError:
             pass
 
@@ -2914,6 +3118,10 @@ class Tools:
         self.secret_grants = set()
         self.read_files = set()
         self.live = False
+        self.diff_preview = None
+
+    def note_diff(self, path, old, new):
+        self.diff_preview = {"path": path, "old": old, "new": new}
 
     def box_engine(self):
         try:
@@ -3113,7 +3321,7 @@ class Tools:
             return "error: no such file"
         return truncate(done.stdout.decode("utf-8", "replace"), 8000, 2000)
 
-    def box_write(self, path, content, old=_UNSET, action="write", message=""):
+    def box_write(self, path, content, old=_UNSET, action="write", message="", preview=True):
         engine, failure = self.box_engine()
         if not engine:
             return failure
@@ -3138,6 +3346,8 @@ class Tools:
         self.read_files.add(("box", path))
         audited = self.audit_file(action, "sandbox:" + path, old, content, message, (), old_size, None)
         note = "" if audited is not False else " (audit failed)"
+        if preview and old_size is None:
+            self.note_diff(path, old, content)
         return "wrote %s%s" % (path, note)
 
     def box_old_image(self, path):
@@ -3172,6 +3382,8 @@ class Tools:
             return "error: " + short_error(exc)
         audited = self.audit_file("write", os.path.abspath(path), old, content, "", (), old_size, None)
         note = "" if audited is not False else " (audit failed)"
+        if old_size is None:
+            self.note_diff(path, old, content)
         return "wrote %s%s" % (path, note)
 
     def run_edit_file(self, args):
@@ -3200,6 +3412,7 @@ class Tools:
             return "error: " + short_error(exc)
         audited = self.audit_file("edit", os.path.abspath(path), content, updated)
         note = "" if audited is not False else " (audit failed)"
+        self.note_diff(path, content, updated)
         return "edited %s%s" % (path, note)
 
     def run_web_search(self, args):
@@ -3864,6 +4077,8 @@ class Tools:
         self.read_files.discard(("home", os.path.abspath(path)))
         audited = self.audit_file("delete", os.path.abspath(path), old, None, "", (), old_size, None)
         note = "" if audited is not False else " (audit failed)"
+        if old_size is None:
+            self.note_diff(path, old, None)
         return "deleted %s%s" % (path, note)
 
     def box_delete(self, path):
@@ -3890,6 +4105,8 @@ class Tools:
         self.read_files.discard(("box", path))
         audited = self.audit_file("delete", "sandbox:" + path, old, None, "", (), old_size, None)
         note = "" if audited is not False else " (audit failed)"
+        if old_size is None:
+            self.note_diff(path, old, None)
         return "deleted %s%s" % (path, note)
 
     def run_audit(self, args):
@@ -3941,6 +4158,10 @@ class Tools:
         message = "restored from audit #%d (%s-image)" % (row_id, which)
         if row["path"].startswith("sandbox:"):
             return self.box_restore(row["path"][len("sandbox:"):], target, message)
+        if os.path.isfile(row["path"]):
+            current, capped = read_capped(row["path"])
+        else:
+            current, capped = None, None
         try:
             parent = os.path.dirname(row["path"])
             if parent:
@@ -3952,12 +4173,17 @@ class Tools:
             return "error: " + short_error(exc)
         self.audit_file("restore", row["path"], None, target, message)
         self.app.store.upsert_file_record(row["path"], target, None, self.active_profile)
+        if capped is None:
+            self.note_diff(row["path"], current, target)
         return "restored %s from audit #%d (%s-image)" % (row["path"], row_id, which)
 
     def box_restore(self, path, target, message):
         self.read_files.add(("box", path))
-        written = self.box_write(path, target, old=None, action="restore", message=message)
+        old, capped = self.box_old_image(path)
+        written = self.box_write(path, target, old=None, action="restore", message=message, preview=False)
         if written.startswith("wrote "):
+            if capped is None:
+                self.note_diff(path, old, target)
             return "restored sandbox:%s (%s)" % (path, message)
         return written
 
@@ -4079,6 +4305,10 @@ class Agent:
         self.timed_out = False
         self.runner_override = None
         self.user_stepped_in = False
+        self._turn_open = False
+        self._turn_goal = ""
+        self._turn_steps = 0
+        self._turn_tools = 0
         self.switch_profile(config.profile, silent=True)
 
     def apply_system(self):
@@ -4200,8 +4430,34 @@ class Agent:
     def show_result(self, result, elapsed_ms):
         if self.quiet:
             return
+        self.show_diff()
         print(paint("└─ %d chars · %dms" % (len(result), elapsed_ms), Ansi.DIM))
 
+    def show_diff(self):
+        preview = getattr(self.tools, "diff_preview", None)
+        if preview is None:
+            return
+        self.tools.diff_preview = None
+        rendered = render_diff(preview["path"], preview["old"], preview["new"], colors=color_enabled())
+        if rendered:
+            print(rendered)
+
+    def checkpoint(self):
+        if not self.persist:
+            return
+        self.store.save_session(self.profile, self.messages, self.bot)
+        if self._turn_open:
+            self.store.save_turn_state(self.profile, self.bot, {"open": True, "goal": self._turn_goal, "steps": self._turn_steps, "tools": self._turn_tools, "updated": now_iso()})
+
+    def close_turn(self):
+        if not self.persist:
+            self._turn_open = False
+            return
+        self.store.save_session(self.profile, self.messages, self.bot)
+        self._turn_open = False
+        if self._turn_goal:
+            self.store.save_turn_state(self.profile, self.bot, {"open": False, "goal": self._turn_goal, "steps": self._turn_steps, "tools": self._turn_tools, "updated": now_iso()})
+
     def run_turn(self, text, capture=False, max_steps=MAX_STEPS, _routed=False):
         if not _routed:
             mentioned = BOT_MENTION_RE.match(text or "")
@@ -4214,6 +4470,11 @@ class Agent:
                 text = mentioned.group(2)
         self.messages.append({"role": "user", "content": text})
         self.store.log_event(self.profile, "user", "message", self.store.redact(text, self.profile))
+        self._turn_open = True
+        self._turn_goal = text[:500]
+        self._turn_steps = 0
+        self._turn_tools = 0
+        self.checkpoint()
         started = time.time()
         backends = []
         pieces = []
@@ -4250,6 +4511,7 @@ class Agent:
                 if loud:
                     print(paint("backend error: %s" % exc, Ansi.RED))
                 self.messages.pop()
+                self.checkpoint()
                 return last_text or "backend error: %s" % exc
             finally:
                 if spinner is not None:
@@ -4316,6 +4578,9 @@ class Agent:
                     last_pair = pair
                     run_len = 1
                     nudged = False
+            self._turn_steps += 1
+            self._turn_tools += len(step_pairs)
+            self.checkpoint()
             if productive:
                 stall = 0
             else:
@@ -4342,8 +4607,7 @@ class Agent:
         else:
             if loud:
                 print(paint("step budget exhausted (%d total steps)" % TOTAL_STEP_CAP, Ansi.YELLOW))
-        if self.persist:
-            self.store.save_session(self.profile, self.messages, self.bot)
+        self.close_turn()
         total = time.time() - started
         if loud:
             print(paint("tokens≈%d · %s · %.1fs" % (conversation_tokens(self.messages), "+".join(backends), total), Ansi.DIM))
@@ -4773,6 +5037,7 @@ def handle_command(agent, text):
             print(paint("use /restore <audit id>", Ansi.RED))
             return True
         print(agent.tools.dispatch("restore", json.dumps({"id": row_id})))
+        agent.show_diff()
     elif name == "release":
         bits = arg.split(None, 1)
         if len(bits) != 2 or bits[0] not in ("major", "minor", "patch"):
@@ -5408,6 +5673,7 @@ def build_parser():
     parser.add_argument("--yes", action="store_true")
     parser.add_argument("--yolo", action="store_true")
     parser.add_argument("--auto", action="store_true")
+    parser.add_argument("--continue", dest="resume", action="store_true")
     parser.add_argument("--env", default=None)
     parser.add_argument("--version", action="store_true")
     parser.add_argument("--telegram", action="store_true")
@@ -5433,9 +5699,22 @@ def banner(agent):
         print("  %s: %s" % (label, probe_backend(models_url, key)))
     if not agent.config.devplace_key:
         print("  fallback: no DEVPLACE_API_KEY")
+    if agent.store.load_turn_state(agent.profile, agent.bot).get("open"):
+        print(paint("interrupted turn available, restart with --continue to pick it up", Ansi.YELLOW))
     print(paint("type /help for commands", Ansi.DIM))
 
 
+def resume_turn(agent):
+    state = agent.store.load_turn_state(agent.profile, agent.bot)
+    if not state.get("open"):
+        print("nothing to resume: last turn for %s/%s completed normally" % (agent.profile, agent.bot))
+        return ""
+    goal = state.get("goal") or ""
+    print("resuming interrupted turn for %s/%s: %d steps, %d tool calls, goal: %s" % (agent.profile, agent.bot, state.get("steps", 0), state.get("tools", 0), goal[:200]))
+    message = "Continue where you left off%s. Do not repeat steps already completed above; carry on with the next step." % (": the interrupted request was: %s" % goal if goal else "")
+    return agent.run_turn(message)
+
+
 def repl(agent):
     if readline is not None:
         try:
@@ -5469,6 +5748,7 @@ def repl(agent):
                 agent.run_turn(text)
             except KeyboardInterrupt:
                 print(paint("\ninterrupted", Ansi.YELLOW))
+                agent.checkpoint()
     finally:
         scheduler_stop.set()
         if readline is not None:
@@ -5526,8 +5806,11 @@ def main(argv=None):
         try:
             agent.run_turn(" ".join(args.prompt))
         finally:
+            agent.checkpoint()
             agent.store.close()
         return 0
+    if args.resume and not args.telegram and not args.scheduler:
+        resume_turn(agent)
     banner(agent)
     try:
         repl(agent)