feat(webui): stream live file edit events

nanobot/utils/file_edit_events.py

@@ -4,13 +4,17 @@ from __future__ import annotations
 
 import difflib
 import json
-from dataclasses import dataclass
+import re
+import time
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any
+from typing import Any, Awaitable, Callable
 
 
 TRACKED_FILE_EDIT_TOOLS = frozenset({"write_file", "edit_file", "notebook_edit"})
 _MAX_SNAPSHOT_BYTES = 2 * 1024 * 1024
+_LIVE_EMIT_INTERVAL_S = 0.18
+_LIVE_EMIT_LINE_STEP = 24
 
 
 @dataclass(slots=True)
@@ -103,6 +107,8 @@ def line_diff_stats(before: str | None, after: str | None) -> tuple[int, int]:
     """Return ``(added, deleted)`` for a UTF-8 text line-level diff."""
     if before is None or after is None:
         return 0, 0
+    if before == "":
+        return _text_line_count(after), 0
     before_lines = before.replace("\r\n", "\n").splitlines()
     after_lines = after.replace("\r\n", "\n").splitlines()
     added = 0
@@ -118,6 +124,28 @@ def line_diff_stats(before: str | None, after: str | None) -> tuple[int, int]:
     return added, deleted
 
 
+def _text_line_count(text: str) -> int:
+    if not text:
+        return 0
+    line_count = 0
+    last_was_newline = False
+    last_was_cr = False
+    for ch in text:
+        if ch == "\r":
+            line_count += 1
+            last_was_newline = True
+            last_was_cr = True
+        elif ch == "\n":
+            if not last_was_cr:
+                line_count += 1
+            last_was_newline = True
+            last_was_cr = False
+        else:
+            last_was_newline = False
+            last_was_cr = False
+    return line_count if last_was_newline else line_count + 1
+
+
 def prepare_file_edit_tracker(
     *,
     call_id: str,
@@ -160,12 +188,22 @@ def build_file_edit_start_event(
     )
 
 
-def build_file_edit_end_event(tracker: FileEditTracker) -> dict[str, Any]:
+def build_file_edit_end_event(
+    tracker: FileEditTracker,
+    params: dict[str, Any] | None = None,
+) -> dict[str, Any]:
     after = read_file_snapshot(tracker.path)
+    counted = False
     if tracker.before.countable and after.countable:
         added, deleted = line_diff_stats(tracker.before.text, after.text)
+        counted = True
     else:
-        added, deleted = 0, 0
+        predicted_after = _predict_after_text(tracker.tool, params or {}, tracker.before)
+        if tracker.before.countable and predicted_after is not None:
+            added, deleted = line_diff_stats(tracker.before.text, predicted_after)
+            counted = True
+        else:
+            added, deleted = 0, 0
     return _event_payload(
         tracker,
         phase="end",
@@ -173,11 +211,14 @@ def build_file_edit_end_event(tracker: FileEditTracker) -> dict[str, Any]:
         added=added,
         deleted=deleted,
         approximate=False,
-        binary=after.binary or after.oversized or after.unreadable,
+        binary=(after.binary or after.oversized or after.unreadable) and not counted,
     )
 
 
-def build_file_edit_error_event(tracker: FileEditTracker, error: str | None = None) -> dict[str, Any]:
+def build_file_edit_error_event(
+    tracker: FileEditTracker,
+    error: str | None = None,
+) -> dict[str, Any]:
     payload = _event_payload(
         tracker,
         phase="error",
@@ -191,6 +232,427 @@ def build_file_edit_error_event(tracker: FileEditTracker, error: str | None = No
     return payload
 
 
+def build_file_edit_live_event(
+    tracker: FileEditTracker,
+    *,
+    added: int,
+    deleted: int = 0,
+) -> dict[str, Any]:
+    """Build an approximate in-progress event while tool-call arguments stream."""
+    return _event_payload(
+        tracker,
+        phase="start",
+        status="editing",
+        added=added,
+        deleted=deleted,
+        approximate=True,
+    )
+
+
+def build_file_edit_pending_event(
+    *,
+    call_id: str,
+    tool_name: str,
+    added: int = 0,
+    deleted: int = 0,
+) -> dict[str, Any]:
+    """Build an early placeholder before the streamed JSON path is available."""
+    return {
+        "version": 1,
+        "call_id": str(call_id or ""),
+        "tool": tool_name,
+        "path": "",
+        "phase": "start",
+        "added": max(0, int(added)),
+        "deleted": max(0, int(deleted)),
+        "approximate": True,
+        "status": "editing",
+        "pending": True,
+    }
+
+
+class StreamingFileEditTracker:
+    """Track file-edit tool arguments while the model is still streaming them.
+
+    Tool execution events only begin after the provider has completed the full
+    function call.  For large ``write_file`` calls, the long wait is usually the
+    model producing the JSON ``content`` argument.  Large ``edit_file`` calls
+    can have the same wait while ``old_text`` / ``new_text`` stream in.  This
+    tracker converts those argument deltas into approximate WebUI file-edit
+    events before the final exact diff is available.
+    """
+
+    def __init__(
+        self,
+        *,
+        workspace: Path | None,
+        tools: Any,
+        emit: Callable[[list[dict[str, Any]]], Awaitable[None]],
+    ) -> None:
+        self._workspace = workspace
+        self._tools = tools
+        self._emit = emit
+        self._states: dict[str, _StreamingFileEditState] = {}
+
+    async def update(self, payload: dict[str, Any]) -> None:
+        key = _stream_key(payload)
+        if not key:
+            return
+        state = self._states.get(key)
+        if state is None:
+            state = _StreamingFileEditState(key=key)
+            self._states[key] = state
+
+        state.apply_delta(payload)
+        if state.name not in {"write_file", "edit_file"}:
+            return
+        if state.path is None:
+            state.path = _extract_complete_json_string(state.arguments, "path")
+        if state.path is None:
+            added, deleted = state.live_diff_counts()
+            now = time.monotonic()
+            if state.should_emit_pending(added, deleted, now):
+                state.mark_pending_emitted(added, deleted, now)
+                await self._emit([build_file_edit_pending_event(
+                    call_id=state.call_id or state.key,
+                    tool_name=state.name,
+                    added=added,
+                    deleted=deleted,
+                )])
+            return
+        if state.tracker is None:
+            tool = self._tools.get(state.name) if hasattr(self._tools, "get") else None
+            state.tracker = prepare_file_edit_tracker(
+                call_id=state.call_id or state.key,
+                tool_name=state.name,
+                tool=tool,
+                workspace=self._workspace,
+                params={"path": state.path},
+            )
+            if state.tracker is None:
+                return
+
+        added, deleted = state.live_diff_counts()
+        now = time.monotonic()
+        if not state.should_emit(added, deleted, now):
+            return
+        state.mark_emitted(added, deleted, now)
+        await self._emit([build_file_edit_live_event(
+            state.tracker,
+            added=added,
+            deleted=deleted,
+        )])
+
+    async def flush(self) -> None:
+        events: list[dict[str, Any]] = []
+        now = time.monotonic()
+        for state in self._states.values():
+            if state.tracker is None:
+                continue
+            added, deleted = state.live_diff_counts()
+            if (
+                state.last_emitted_added == added
+                and state.last_emitted_deleted == deleted
+                and state.emitted_once
+            ):
+                continue
+            state.mark_emitted(added, deleted, now)
+            events.append(build_file_edit_live_event(
+                state.tracker,
+                added=added,
+                deleted=deleted,
+            ))
+        if events:
+            await self._emit(events)
+
+    def apply_final_call_ids(self, final_tool_calls: list[Any]) -> None:
+        """Keep final start/end events keyed to any earlier streamed placeholder."""
+        for tool_call in final_tool_calls:
+            canonical = self.canonical_call_id_for(tool_call)
+            if canonical:
+                try:
+                    tool_call.id = canonical
+                except Exception:
+                    pass
+
+    def canonical_call_id_for(self, tool_call: Any) -> str | None:
+        for state in self._states.values():
+            if state.matches_final_tool_call(tool_call):
+                return state.call_id or (state.tracker.call_id if state.tracker else None) or state.key
+        return None
+
+    async def error_unmatched(
+        self,
+        final_tool_calls: list[Any],
+        error: str,
+    ) -> None:
+        """Mark streamed edits as failed when no final tool call will run."""
+        events: list[dict[str, Any]] = []
+        for state in self._states.values():
+            if state.tracker is None:
+                continue
+            if any(state.matches_final_tool_call(tool_call) for tool_call in final_tool_calls):
+                continue
+            events.append(build_file_edit_error_event(state.tracker, error))
+        if events:
+            await self._emit(events)
+
+
+@dataclass(slots=True)
+class _StreamingJsonStringField:
+    key: str
+    scan_pos: int | None = None
+    closed: bool = False
+    escape: bool = False
+    unicode_remaining: int = 0
+    unicode_buffer: str = ""
+    newline_count: int = 0
+    has_chars: bool = False
+    last_char_newline: bool = False
+    last_char_cr: bool = False
+
+    @property
+    def line_count(self) -> int:
+        if not self.has_chars:
+            return 0
+        return self.newline_count + (0 if self.last_char_newline else 1)
+
+    def reset(self) -> None:
+        self.scan_pos = None
+        self.closed = False
+        self.escape = False
+        self.unicode_remaining = 0
+        self.unicode_buffer = ""
+        self.newline_count = 0
+        self.has_chars = False
+        self.last_char_newline = False
+        self.last_char_cr = False
+
+    def scan(self, source: str) -> None:
+        if self.closed:
+            return
+        if self.scan_pos is None:
+            match = re.search(rf'"{re.escape(self.key)}"\s*:\s*"', source)
+            if match is None:
+                return
+            self.scan_pos = match.end()
+        i = self.scan_pos
+        while i < len(source):
+            ch = source[i]
+            if self.unicode_remaining > 0:
+                self.unicode_buffer += ch
+                self.unicode_remaining -= 1
+                if self.unicode_remaining == 0:
+                    try:
+                        decoded = chr(int(self.unicode_buffer, 16))
+                    except ValueError:
+                        decoded = "x"
+                    self.unicode_buffer = ""
+                    self._mark_char(decoded)
+                i += 1
+                continue
+            if self.escape:
+                self.escape = False
+                if ch == "u":
+                    self.unicode_remaining = 4
+                    self.unicode_buffer = ""
+                elif ch == "n":
+                    self._mark_char("\n")
+                elif ch == "r":
+                    self._mark_char("\r")
+                else:
+                    self._mark_char(ch)
+                i += 1
+                continue
+            if ch == "\\":
+                self.escape = True
+                i += 1
+                continue
+            if ch == '"':
+                self.closed = True
+                i += 1
+                break
+            self._mark_char(ch)
+            i += 1
+        self.scan_pos = i
+
+    def _mark_char(self, ch: str) -> None:
+        self.has_chars = True
+        if ch == "\r":
+            self.newline_count += 1
+            self.last_char_newline = True
+            self.last_char_cr = True
+        elif ch == "\n":
+            if not self.last_char_cr:
+                self.newline_count += 1
+            self.last_char_newline = True
+            self.last_char_cr = False
+        else:
+            self.last_char_newline = False
+            self.last_char_cr = False
+
+
+@dataclass(slots=True)
+class _StreamingFileEditState:
+    key: str
+    call_id: str = ""
+    name: str = ""
+    arguments: str = ""
+    path: str | None = None
+    tracker: FileEditTracker | None = None
+    content: _StreamingJsonStringField = field(
+        default_factory=lambda: _StreamingJsonStringField("content")
+    )
+    old_text: _StreamingJsonStringField = field(
+        default_factory=lambda: _StreamingJsonStringField("old_text")
+    )
+    new_text: _StreamingJsonStringField = field(
+        default_factory=lambda: _StreamingJsonStringField("new_text")
+    )
+    emitted_once: bool = False
+    last_emitted_added: int = -1
+    last_emitted_deleted: int = -1
+    last_emit_at: float = 0.0
+    pending_emitted: bool = False
+    last_pending_added: int = -1
+    last_pending_deleted: int = -1
+    last_pending_at: float = 0.0
+
+    def apply_delta(self, payload: dict[str, Any]) -> None:
+        call_id = payload.get("call_id")
+        if isinstance(call_id, str) and call_id:
+            self.call_id = call_id
+        name = payload.get("name")
+        if isinstance(name, str) and name:
+            self.name = name
+        args = payload.get("arguments")
+        if isinstance(args, str):
+            self.arguments = args
+            self.content.reset()
+            self.old_text.reset()
+            self.new_text.reset()
+            return
+        delta = payload.get("arguments_delta")
+        if isinstance(delta, str) and delta:
+            self.arguments += delta
+
+    def live_diff_counts(self) -> tuple[int, int]:
+        if self.name == "write_file":
+            self.content.scan(self.arguments)
+            return self.content.line_count, 0
+        if self.name == "edit_file":
+            self.old_text.scan(self.arguments)
+            self.new_text.scan(self.arguments)
+            return self.new_text.line_count, self.old_text.line_count
+        return 0, 0
+
+    def should_emit(self, added: int, deleted: int, now: float) -> bool:
+        if not self.emitted_once:
+            return True
+        if added == self.last_emitted_added and deleted == self.last_emitted_deleted:
+            return False
+        if max(
+            abs(added - self.last_emitted_added),
+            abs(deleted - self.last_emitted_deleted),
+        ) >= _LIVE_EMIT_LINE_STEP:
+            return True
+        return now - self.last_emit_at >= _LIVE_EMIT_INTERVAL_S
+
+    def mark_emitted(self, added: int, deleted: int, now: float) -> None:
+        self.emitted_once = True
+        self.last_emitted_added = added
+        self.last_emitted_deleted = deleted
+        self.last_emit_at = now
+
+    def should_emit_pending(self, added: int, deleted: int, now: float) -> bool:
+        if not self.pending_emitted:
+            return True
+        if added == self.last_pending_added and deleted == self.last_pending_deleted:
+            return False
+        if max(
+            abs(added - self.last_pending_added),
+            abs(deleted - self.last_pending_deleted),
+        ) >= _LIVE_EMIT_LINE_STEP:
+            return True
+        return now - self.last_pending_at >= _LIVE_EMIT_INTERVAL_S
+
+    def mark_pending_emitted(self, added: int, deleted: int, now: float) -> None:
+        self.pending_emitted = True
+        self.last_pending_added = added
+        self.last_pending_deleted = deleted
+        self.last_pending_at = now
+
+    def matches_final_tool_call(self, tool_call: Any) -> bool:
+        call_id = getattr(tool_call, "id", None)
+        canonical = self.call_id or (self.tracker.call_id if self.tracker else "")
+        if isinstance(call_id, str) and call_id and canonical and call_id == canonical:
+            return True
+        name = getattr(tool_call, "name", None)
+        if name != self.name:
+            return False
+        arguments = getattr(tool_call, "arguments", None)
+        if not isinstance(arguments, dict):
+            return False
+        path = arguments.get("path")
+        if self.path is None and isinstance(path, str) and path:
+            self.path = path
+            return True
+        return isinstance(path, str) and path == self.path
+
+
+def _stream_key(payload: dict[str, Any]) -> str:
+    index = payload.get("index")
+    if isinstance(index, int):
+        return f"idx:{index}"
+    if isinstance(index, str) and index:
+        return f"idx:{index}"
+    call_id = payload.get("call_id")
+    if isinstance(call_id, str) and call_id:
+        return f"id:{call_id}"
+    return ""
+
+
+def _extract_complete_json_string(source: str, key: str) -> str | None:
+    match = re.search(rf'"{re.escape(key)}"\s*:\s*"', source)
+    if match is None:
+        return None
+    out: list[str] = []
+    i = match.end()
+    escape = False
+    while i < len(source):
+        ch = source[i]
+        if escape:
+            escape = False
+            if ch == "n":
+                out.append("\n")
+            elif ch == "r":
+                out.append("\r")
+            elif ch == "t":
+                out.append("\t")
+            elif ch == "u":
+                digits = source[i + 1:i + 5]
+                if len(digits) < 4:
+                    return None
+                try:
+                    out.append(chr(int(digits, 16)))
+                except ValueError:
+                    return None
+                i += 4
+            else:
+                out.append(ch)
+            i += 1
+            continue
+        if ch == "\\":
+            escape = True
+            i += 1
+            continue
+        if ch == '"':
+            return "".join(out)
+        out.append(ch)
+        i += 1
+    return None
+
+
 def _event_payload(
     tracker: FileEditTracker,
     *,
@@ -206,6 +668,7 @@ def _event_payload(
         "call_id": tracker.call_id,
         "tool": tracker.tool,
         "path": tracker.display_path,
+        "absolute_path": tracker.path.as_posix(),
         "phase": phase,
         "added": max(0, int(added)),
         "deleted": max(0, int(deleted)),
@@ -260,8 +723,14 @@ def _predict_notebook_after_text(params: dict[str, Any], before_text: str) -> st
         return None
     new_source = params.get("new_source")
     source = new_source if isinstance(new_source, str) else ""
-    cell_type = params.get("cell_type") if params.get("cell_type") in ("code", "markdown") else "code"
-    mode = params.get("edit_mode") if params.get("edit_mode") in ("replace", "insert", "delete") else "replace"
+    cell_type = (
+        params.get("cell_type") if params.get("cell_type") in ("code", "markdown") else "code"
+    )
+    mode = (
+        params.get("edit_mode")
+        if params.get("edit_mode") in ("replace", "insert", "delete")
+        else "replace"
+    )
     if mode == "delete":
         if 0 <= cell_index < len(cells):
             cells.pop(cell_index)

nanobot/utils/webui_transcript.py

@@ -144,6 +144,17 @@ def replay_transcript_to_ui_messages(
     def _ensure_activity_segment() -> str:
         return active_activity_segment_id or _new_activity_segment()
 
+    def close_activity_for_answer() -> None:
+        nonlocal active_activity_segment_id, active_file_edit_segment_id
+        active_activity_segment_id = None
+        active_file_edit_segment_id = None
+
+    def close_file_edit_phase_before_activity() -> None:
+        nonlocal active_activity_segment_id, active_file_edit_segment_id
+        if active_file_edit_segment_id:
+            active_activity_segment_id = None
+            active_file_edit_segment_id = None
+
     def attach_reasoning_chunk(prev: list[dict[str, Any]], chunk: str, idx: int) -> None:
         for i in range(len(prev) - 1, -1, -1):
             candidate = prev[i]
@@ -243,7 +254,7 @@ def replay_transcript_to_ui_messages(
                 return
 
     def absorb_complete(extra: dict[str, Any], idx: int) -> None:
-        nonlocal active_activity_segment_id
+        nonlocal active_activity_segment_id, active_file_edit_segment_id
         last = messages[-1] if messages else None
         if last and is_reasoning_only_placeholder(last):
             messages[-1] = {
@@ -262,35 +273,50 @@ def replay_transcript_to_ui_messages(
                 },
             )
         active_activity_segment_id = None
+        active_file_edit_segment_id = None
 
     def _file_edit_key(edit: dict[str, Any]) -> str:
-        return "|".join(
-            str(edit.get(k) or "")
-            for k in ("call_id", "tool", "path")
-        )
+        call_id = str(edit.get("call_id") or "")
+        tool = str(edit.get("tool") or "")
+        if call_id:
+            return f"{call_id}|{tool}"
+        return f"{tool}|{edit.get('path') or ''}"
+
+    def find_file_edit_trace_index(
+        segment: str | None,
+        edits: list[dict[str, Any]],
+    ) -> int | None:
+        incoming_keys = {_file_edit_key(edit) for edit in edits if isinstance(edit, dict)}
+        for i in range(len(messages) - 1, -1, -1):
+            candidate = messages[i]
+            if candidate.get("role") == "user":
+                break
+            if candidate.get("kind") != "trace" or not candidate.get("fileEdits"):
+                continue
+            if segment and candidate.get("activitySegmentId") == segment:
+                return i
+            existing_edits = candidate.get("fileEdits")
+            if not isinstance(existing_edits, list):
+                continue
+            for existing in existing_edits:
+                if isinstance(existing, dict) and _file_edit_key(existing) in incoming_keys:
+                    return i
+        return None
 
     def upsert_file_edits(edits: list[dict[str, Any]], idx: int) -> None:
         nonlocal active_file_edit_segment_id
         if not edits:
             return
-        last = messages[-1] if messages else None
-        if (
-            active_file_edit_segment_id
-            and last
-            and last.get("kind") == "trace"
-            and last.get("fileEdits")
-        ):
-            segment = active_file_edit_segment_id
+        segment = active_file_edit_segment_id
+        target_index = find_file_edit_trace_index(segment, edits)
+        if target_index is not None:
+            last = messages[target_index]
+            segment = str(last.get("activitySegmentId") or segment or _new_activity_segment(activate=False))
+            active_file_edit_segment_id = segment
         else:
-            segment = _new_activity_segment(activate=False)
+            if not segment:
+                segment = _new_activity_segment(activate=False)
             active_file_edit_segment_id = segment
-        if not (
-            last
-            and last.get("kind") == "trace"
-            and not last.get("isStreaming")
-            and last.get("fileEdits")
-            and last.get("activitySegmentId") == segment
-        ):
             messages.append(
                 {
                     "id": _new_id("tr", idx),
@@ -303,7 +329,11 @@ def replay_transcript_to_ui_messages(
                     "createdAt": _ts_base + idx,
                 },
             )
-            last = messages[-1]
+            target_index = len(messages) - 1
+            last = messages[target_index]
+        if not segment:
+            segment = _new_activity_segment(activate=False)
+            active_file_edit_segment_id = segment
         existing = list(last.get("fileEdits") or [])
         index_by_key = {
             _file_edit_key(edit): pos
@@ -316,11 +346,14 @@ def replay_transcript_to_ui_messages(
             key = _file_edit_key(edit)
             if key in index_by_key:
                 pos = index_by_key[key]
-                existing[pos] = {**existing[pos], **edit}
+                merged = {**existing[pos], **edit}
+                if edit.get("path") and not edit.get("pending"):
+                    merged.pop("pending", None)
+                existing[pos] = merged
             else:
                 index_by_key[key] = len(existing)
                 existing.append(dict(edit))
-        messages[-1] = {
+        messages[target_index] = {
             **last,
             "fileEdits": existing,
             "activitySegmentId": last.get("activitySegmentId") or segment,
@@ -365,6 +398,7 @@ def replay_transcript_to_ui_messages(
             chunk = rec.get("text")
             if not isinstance(chunk, str):
                 continue
+            close_activity_for_answer()
             adopted = find_active_placeholder(messages) if buffer_message_id is None else None
             if buffer_message_id is None:
                 if adopted:
@@ -403,6 +437,7 @@ def replay_transcript_to_ui_messages(
             chunk = rec.get("text")
             if not isinstance(chunk, str) or not chunk:
                 continue
+            close_file_edit_phase_before_activity()
             attach_reasoning_chunk(messages, chunk, idx)
             continue
 
@@ -424,6 +459,7 @@ def replay_transcript_to_ui_messages(
                 line = rec.get("text")
                 if not isinstance(line, str) or not line:
                     continue
+                close_file_edit_phase_before_activity()
                 attach_reasoning_chunk(messages, line, idx)
                 close_reasoning(messages)
                 continue