better agent

cht/ui/agent_output.py

@@ -1,20 +1,35 @@
-"""Agent output panel — scrollable text view with markdown rendering."""
+"""Agent output panel — single TextView, fully selectable/copy-pastable."""
+
+import logging
 
 import gi
 gi.require_version("Gtk", "4.0")
-from gi.repository import Gtk
+from gi.repository import Gtk, GLib, Pango
 
 from cht.ui import markdown
+from cht.agent.base import (
+    AssistantMessage,
+    ImageBlock,
+    TextBlock,
+    Thread,
+    ToolResult,
+    ToolUse,
+    TranscriptBlock,
+    UserMessage,
+)
+
+log = logging.getLogger(__name__)
 
 
 class AgentOutputPanel(Gtk.Frame):
-    """Scrollable text view that displays agent responses with markdown."""
+    """Scrollable text view showing the full conversation, copy-pastable."""
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
         box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=0)
 
+        # Header
         header = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=4)
         header.set_margin_start(8)
         header.set_margin_end(8)
@@ -33,13 +48,16 @@ class AgentOutputPanel(Gtk.Frame):
 
         box.append(header)
 
+        # Single text view for the whole conversation
         self._view = Gtk.TextView()
         self._view.set_editable(False)
         self._view.set_wrap_mode(Gtk.WrapMode.WORD_CHAR)
         self._view.set_cursor_visible(False)
         self._view.set_left_margin(8)
         self._view.set_right_margin(8)
-        markdown.setup_tags(self._view.get_buffer())
+        self._view.set_top_margin(4)
+        self._view.set_bottom_margin(4)
+        self._setup_tags()
 
         scroll = Gtk.ScrolledWindow()
         scroll.set_vexpand(True)
@@ -48,52 +66,170 @@ class AgentOutputPanel(Gtk.Frame):
 
         self.set_child(box)
 
-        # Streaming state
-        self._thinking_replaced = False
-        self._response_start_mark = None
-        self._response_accum = []
+        # Streaming state: track where the current assistant response starts
+        self._response_marks: dict[str, Gtk.TextMark] = {}
+        self._response_accum: dict[str, list[str]] = {}
+
+    def _setup_tags(self):
+        buf = self._view.get_buffer()
+        markdown.setup_tags(buf)
+        buf.create_tag("user-prefix", weight=700, foreground="#7aafff")
+        buf.create_tag("assistant-prefix", weight=700, foreground="#8bc78b")
+        buf.create_tag("tool-prefix", weight=700, foreground="#d4a053")
+        buf.create_tag("tool-output", foreground="#aaaaaa", left_margin=16)
+        buf.create_tag("error", foreground="#ff6b6b")
+        buf.create_tag("ref-chip", foreground="#7aafff", style=Pango.Style.ITALIC)
+        buf.create_tag("status", foreground="#888888")
+
+    # -- Public API --
 
     def append(self, text: str) -> None:
-        """Append plain text and auto-scroll."""
+        """Append a status/info line."""
         buf = self._view.get_buffer()
-        buf.insert(buf.get_end_iter(), text)
-        self._view.scroll_to_iter(buf.get_end_iter(), 0, False, 0, 0)
+        end = buf.get_end_iter()
+        mark = buf.create_mark(None, end, True)
+        buf.insert(end, text)
+        start = buf.get_iter_at_mark(mark)
+        buf.apply_tag_by_name("status", start, buf.get_end_iter())
+        buf.delete_mark(mark)
+        self._scroll_to_bottom()
 
     def clear(self) -> None:
-        """Clear all output."""
         self._view.get_buffer().set_text("")
+        self._response_marks.clear()
+        self._response_accum.clear()
 
-    def begin_response(self) -> None:
-        """Reset streaming state for a new response (call before sending)."""
-        self._thinking_replaced = False
-        self._response_start_mark = None
-        self._response_accum = []
-
-    def replace_thinking(self, chunk: str) -> None:
-        """Replace the '...' placeholder with streamed chunks."""
+    def add_user_message(self, text: str, frames: list | None = None,
+                         transcripts: list | None = None) -> None:
         buf = self._view.get_buffer()
-        if not self._thinking_replaced:
-            self._thinking_replaced = True
-            end = buf.get_end_iter()
-            start = end.copy()
-            start.backward_chars(2)
-            buf.delete(start, end)
-            self._response_start_mark = buf.create_mark(
-                None, buf.get_end_iter(), left_gravity=True
-            )
-        self._response_accum.append(chunk)
-        self.append(chunk)
+        end = buf.get_end_iter()
 
-    def finish_response(self, err: str | None) -> None:
-        """Finalize the response — render markdown or show error."""
-        if err:
-            self.append(f"[Error: {err}]\n")
-            return
-        if self._response_start_mark and self._response_accum:
-            buf = self._view.get_buffer()
-            start = buf.get_iter_at_mark(self._response_start_mark)
+        # Prefix
+        mark = buf.create_mark(None, end, True)
+        buf.insert(end, "\n> ")
+        buf.apply_tag_by_name("user-prefix", buf.get_iter_at_mark(mark), buf.get_end_iter())
+        buf.delete_mark(mark)
+
+        # Text
+        buf.insert(buf.get_end_iter(), text)
+
+        # Reference chips
+        refs = []
+        if frames:
+            for f in frames:
+                fid = f.frame_id if hasattr(f, "frame_id") else (f.id if hasattr(f, "id") else str(f))
+                refs.append(fid)
+        if transcripts:
+            for t in transcripts:
+                tid = t.transcript_id if hasattr(t, "transcript_id") else (t.id if hasattr(t, "id") else str(t))
+                refs.append(tid)
+        if refs:
             end = buf.get_end_iter()
-            buf.delete(start, end)
-            markdown.render(buf, start, "".join(self._response_accum))
-            buf.delete_mark(self._response_start_mark)
-        self.append("\n")
+            mark = buf.create_mark(None, end, True)
+            buf.insert(end, " [" + ", ".join(refs) + "]")
+            buf.apply_tag_by_name("ref-chip", buf.get_iter_at_mark(mark), buf.get_end_iter())
+            buf.delete_mark(mark)
+
+        buf.insert(buf.get_end_iter(), "\n")
+        self._scroll_to_bottom()
+
+    def begin_assistant_message(self, msg_id: str) -> None:
+        buf = self._view.get_buffer()
+        end = buf.get_end_iter()
+
+        # Mark where this response starts (for markdown re-render on finish)
+        self._response_marks[msg_id] = buf.create_mark(f"resp_{msg_id}", end, True)
+        self._response_accum[msg_id] = []
+
+    def append_to_assistant(self, msg_id: str, text: str) -> None:
+        if msg_id not in self._response_accum:
+            return
+        self._response_accum[msg_id].append(text)
+        buf = self._view.get_buffer()
+        buf.insert(buf.get_end_iter(), text)
+        self._scroll_to_bottom()
+
+    def finish_assistant(self, msg_id: str, full_text: str) -> None:
+        mark = self._response_marks.pop(msg_id, None)
+        self._response_accum.pop(msg_id, None)
+        if not mark:
+            return
+        buf = self._view.get_buffer()
+        start = buf.get_iter_at_mark(mark)
+        end = buf.get_end_iter()
+        buf.delete(start, end)
+        it = buf.get_iter_at_mark(mark)
+        markdown.render(buf, it, full_text)
+        buf.insert(buf.get_end_iter(), "\n")
+        buf.delete_mark(mark)
+
+    def add_tool_call(self, tool_use: ToolUse) -> None:
+        buf = self._view.get_buffer()
+        end = buf.get_end_iter()
+
+        mark = buf.create_mark(None, end, True)
+        buf.insert(end, f"  ▶ {tool_use.tool_name}")
+        buf.apply_tag_by_name("tool-prefix", buf.get_iter_at_mark(mark), buf.get_end_iter())
+        buf.delete_mark(mark)
+
+        if tool_use.input:
+            inp = str(tool_use.input)
+            if len(inp) > 80:
+                inp = inp[:77] + "..."
+            end = buf.get_end_iter()
+            mark = buf.create_mark(None, end, True)
+            buf.insert(end, f" {inp}")
+            buf.apply_tag_by_name("tool-output", buf.get_iter_at_mark(mark), buf.get_end_iter())
+            buf.delete_mark(mark)
+
+        buf.insert(buf.get_end_iter(), "\n")
+        self._scroll_to_bottom()
+
+    def update_tool_result(self, tool_use_id: str, result: ToolResult) -> None:
+        buf = self._view.get_buffer()
+        text = result.error or result.output or ""
+        if not text:
+            return
+        end = buf.get_end_iter()
+        mark = buf.create_mark(None, end, True)
+        # Indent tool output
+        indented = "\n".join(f"    {line}" for line in text.split("\n"))
+        tag = "error" if result.error else "tool-output"
+        buf.insert(end, indented + "\n")
+        buf.apply_tag_by_name(tag, buf.get_iter_at_mark(mark), buf.get_end_iter())
+        buf.delete_mark(mark)
+        self._scroll_to_bottom()
+
+    def load_thread(self, thread: Thread) -> None:
+        """Replay a thread to rebuild the conversation view."""
+        self.clear()
+        for msg in thread.messages:
+            if isinstance(msg, UserMessage):
+                text_parts = []
+                frames = []
+                transcripts = []
+                for b in msg.content:
+                    if isinstance(b, TextBlock):
+                        text_parts.append(b.text)
+                    elif isinstance(b, ImageBlock):
+                        frames.append(b)
+                    elif isinstance(b, TranscriptBlock):
+                        transcripts.append(b)
+                self.add_user_message(" ".join(text_parts), frames, transcripts)
+            elif isinstance(msg, AssistantMessage):
+                text = " ".join(b.text for b in msg.content if isinstance(b, TextBlock))
+                buf = self._view.get_buffer()
+                it = buf.get_end_iter()
+                markdown.render(buf, it, text)
+                buf.insert(buf.get_end_iter(), "\n")
+            elif isinstance(msg, ToolUse):
+                self.add_tool_call(msg)
+            elif isinstance(msg, ToolResult):
+                self.update_tool_result(msg.tool_use_id, msg)
+
+    def _scroll_to_bottom(self):
+        def _do():
+            adj = self._view.get_parent().get_vadjustment()
+            adj.set_value(adj.get_upper() - adj.get_page_size())
+            return False
+        GLib.idle_add(_do)