From 84249d5a74321deca3d37d2994473e08c5f5df28 Mon Sep 17 00:00:00 2001
From: "Danilo M." <danix@danix.xyz>
Date: Mon, 24 Aug 2026 10:27:56 +0200
Subject: feat: wrap code blocks and offer a copy link on each

Two things made a fenced code block awkward to read. Qt's markdown
stylesheet wraps p and li but not pre, and that stylesheet is dropped
when only the body fragment is kept, so a long line widened the whole
transcript and left it scrolling sideways. The pre style now carries
white-space:pre-wrap itself.

Copying a block meant selecting it by hand. Each block now gets a copy
link above it, on the right, in the same scheme-dispatch the thinking
and search toggles already use, since Qt rich text has no widgets to
put a real button in. The block text is read back from the source
markdown because Qt emits one <pre> per line, which makes the block
unrecoverable from the rendered HTML; _close_fences runs first, so a
block still streaming is copyable mid-reply. Like the other schemes,
a forged copy link in a reply is defused.
---
 llamachat/ui.py   | 78 +++++++++++++++++++++++++++++++++++++++++++++++++-----
 test_llamachat.py | 79 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 151 insertions(+), 6 deletions(-)

diff --git a/llamachat/ui.py b/llamachat/ui.py
index 3ea3598..7c2e113 100644
--- a/llamachat/ui.py
+++ b/llamachat/ui.py
@@ -16,6 +16,7 @@
 import dataclasses
 import datetime
 import html
+import itertools
 import json
 import re
 from pathlib import Path
@@ -27,8 +28,8 @@ from PySide6.QtGui import (
     QAction, QColor, QDesktopServices, QKeySequence, QPainter, QTextDocument,
 )
 from PySide6.QtWidgets import (
-    QButtonGroup, QComboBox, QDialog, QDialogButtonBox, QFileDialog,
-    QHBoxLayout, QInputDialog, QLabel, QLineEdit, QListWidget,
+    QApplication, QButtonGroup, QComboBox, QDialog, QDialogButtonBox,
+    QFileDialog, QHBoxLayout, QInputDialog, QLabel, QLineEdit, QListWidget,
     QListWidgetItem, QMainWindow, QMenu, QMessageBox, QPlainTextEdit,
     QPushButton, QRadioButton, QSplitter, QTextBrowser, QToolButton,
     QVBoxLayout, QWidget,
@@ -49,6 +50,7 @@ MODE_CHAT = "chat"
 # single message expand independently.
 REASONING_SCHEME = "x-llamachat-reasoning:"
 SEARCH_SCHEME = "x-llamachat-search:"
+COPY_SCHEME = "x-llamachat-copy:"
 
 
 class ContextMeter(QWidget):
@@ -1564,9 +1566,23 @@ class ChatWindow(QMainWindow):
         if target.startswith(SEARCH_SCHEME):
             self._toggle(self.expanded_searches, target[len(SEARCH_SCHEME) :])
             return
+        if target.startswith(COPY_SCHEME):
+            self._copy_code(target[len(COPY_SCHEME) :])
+            return
         if url.scheme() in ("http", "https", "mailto"):
             QDesktopServices.openUrl(url)
 
+    def _copy_code(self, raw: str) -> None:
+        """Copy one code block of one bubble to the clipboard."""
+        bubble_index, _, block_index = raw.partition(".")
+        try:
+            bubble = self.bubbles[int(bubble_index)]
+            block = _code_blocks(_close_fences(bubble["text"]))[int(block_index)]
+        except (ValueError, IndexError, KeyError):
+            return  # a stale link from a re-render; nothing to copy
+        QApplication.clipboard().setText(block)
+        self.show_status("Code block copied.")
+
     def _toggle(self, which: set[int], raw: str) -> None:
         try:
             index = int(raw)
@@ -1743,7 +1759,10 @@ MARKDOWN_FLAGS = (
 # own. Giving those a tint is what makes a code block read as a block.
 _PRE_STYLE = (
     "background:rgba(127,127,127,0.18);padding:1px 6px;margin:0;"
-    "font-family:monospace"
+    # Qt's own markdown stylesheet wraps p and li but not pre, and that
+    # stylesheet is dropped when only the body fragment is kept. Without
+    # this a long line widens the whole transcript instead of wrapping.
+    "font-family:monospace;white-space:pre-wrap"
 )
 
 
@@ -1774,7 +1793,26 @@ def _close_fences(text: str) -> str:
     return text
 
 
-def _markdown_to_fragment(text: str) -> str:
+_BLOCK_RE = re.compile(
+    r"^(?P<indent>\s{0,3})(?P<fence>`{3,}|~{3,})[^\n]*\n"
+    r"(?P<body>.*?)"
+    r"^\s{0,3}(?P=fence)`*~*[ \t]*$",
+    re.MULTILINE | re.DOTALL,
+)
+
+
+def _code_blocks(text: str) -> list[str]:
+    """The text of each fenced block, in the order they appear.
+
+    Copying needs the source, not the rendered HTML: Qt emits one <pre> per
+    line, so the block is not recoverable from the fragment. Parsing the
+    same markdown the renderer sees keeps the two in step, and _close_fences
+    has already terminated a block still streaming.
+    """
+    return [m.group("body").rstrip("\n") for m in _BLOCK_RE.finditer(text)]
+
+
+def _markdown_to_fragment(text: str, index: int | None = None) -> str:
     """Render markdown to an HTML fragment safe to splice into the page.
 
     QTextDocument does the parsing, so there is no markdown dependency. Its
@@ -1801,10 +1839,38 @@ def _markdown_to_fragment(text: str) -> str:
     # so their scheme is defused too.
     fragment = fragment.replace(f'href="{REASONING_SCHEME}', 'href="blocked:')
     fragment = fragment.replace(f'href="{SEARCH_SCHEME}', 'href="blocked:')
+    fragment = fragment.replace(f'href="{COPY_SCHEME}', 'href="blocked:')
 
     # Tint code blocks, which Qt leaves unstyled. Qt emits one <pre> per
     # line, so the tint has to land on every one of them to read as a block.
-    return _PRE_RE.sub(_tint_pre, fragment)
+    fragment = _PRE_RE.sub(_tint_pre, fragment)
+    if index is None:
+        return fragment
+    return _add_copy_links(fragment, index)
+
+
+# The first <pre> of a run is where a copy header goes; the rest are the
+# same block's later lines. A run ends at the first tag that is not a <pre>.
+_PRE_RUN_RE = re.compile(r"(?:<pre\b[^>]*>.*?</pre>\s*)+", re.DOTALL)
+
+
+def _add_copy_links(fragment: str, index: int) -> str:
+    """Put a copy link above each code block.
+
+    Qt rich text has no widgets, so the button is an anchor in the existing
+    scheme-dispatch, the same mechanism the thinking and search toggles use.
+    """
+    counter = itertools.count()
+
+    def header(match: re.Match) -> str:
+        link = (
+            f'<div style="margin:4px 0 0 0;text-align:right">'
+            f'<a href="{COPY_SCHEME}{index}.{next(counter)}" '
+            f'style="color:palette(mid);text-decoration:none">⧉ copy</a></div>'
+        )
+        return link + match.group(0)
+
+    return _PRE_RUN_RE.sub(header, fragment)
 
 
 _PRE_RE = re.compile(r"<pre(\s+style=\"([^\"]*)\")?\s*>")
@@ -1967,7 +2033,7 @@ def _bubble_html(
     # Replies are markdown; what the user typed is shown as typed, so an
     # attached file's contents cannot be reflowed into headings and lists.
     if role == "assistant":
-        body = _markdown_to_fragment(text)
+        body = _markdown_to_fragment(text, index)
     else:
         body = html.escape(text).replace("\n", "<br>")
 
diff --git a/test_llamachat.py b/test_llamachat.py
index bb5f097..7a84608 100755
--- a/test_llamachat.py
+++ b/test_llamachat.py
@@ -441,6 +441,83 @@ def test_unbalanced_backticks():
     print("ok  unbalanced backticks")
 
 
+def test_code_block_wrapping():
+    """A long code line wraps instead of widening the whole transcript."""
+    import os
+
+    os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
+    from PySide6.QtGui import QTextDocument
+    from PySide6.QtWidgets import QApplication
+
+    from llamachat.ui import _markdown_to_fragment
+
+    app = QApplication.instance() or QApplication([])
+    assert app is not None
+
+    # Qt's markdown stylesheet wraps p and li but not pre, and it is dropped
+    # when only the body fragment is kept, so pre has to say so itself.
+    fragment = _markdown_to_fragment("```\n" + "A" * 200 + "\n```\n")
+    assert "white-space:pre-wrap" in fragment, fragment
+
+    # The real check is the layout: no block may exceed the viewport.
+    doc = QTextDocument()
+    doc.setHtml(fragment)
+    doc.setTextWidth(300)
+    layout = doc.documentLayout()
+    widest = 0.0
+    block = doc.begin()
+    while block.isValid():
+        widest = max(widest, layout.blockBoundingRect(block).width())
+        block = block.next()
+    assert widest <= 300, widest
+    print("ok  code block wrapping")
+
+
+def test_code_block_copy_links():
+    """Each fenced block gets a copy link addressing it by index."""
+    import os
+    import re as _re
+
+    os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
+    from PySide6.QtWidgets import QApplication
+
+    from llamachat.ui import (
+        COPY_SCHEME, _close_fences, _code_blocks, _markdown_to_fragment,
+    )
+
+    app = QApplication.instance() or QApplication([])
+    assert app is not None
+
+    # Blocks are read from the source: Qt emits one <pre> per line, so the
+    # text is not recoverable from the rendered fragment.
+    assert _code_blocks("```py\nx=1\ny=2\n```\n") == ["x=1\ny=2"]
+    assert _code_blocks("~~~\na\n~~~\n") == ["a"]
+    assert _code_blocks("no code here") == []
+    # A longer fence can hold shorter ones without ending the block.
+    assert _code_blocks("````\n```\ninner\n```\n````\n") == ["```\ninner\n```"]
+
+    # A link per block, numbered bubble.block so several replies coexist.
+    markdown = "one\n\n```\nAAA\n```\n\ntwo\n\n```\nBBB\n```\n"
+    fragment = _markdown_to_fragment(markdown, 3)
+    assert _re.findall(r"x-llamachat-copy:([0-9.]+)", fragment) == ["3.0", "3.1"]
+
+    # The indices the links carry must select the blocks the parser found.
+    blocks = _code_blocks(_close_fences(markdown))
+    assert blocks == ["AAA", "BBB"], blocks
+
+    # A block still streaming is closed first, so it is copyable mid-reply.
+    assert _code_blocks(_close_fences("```\nhalf")) == ["half"]
+
+    # Without an index there are no links, so a bubble that cannot be
+    # addressed does not emit a link that would not resolve.
+    assert COPY_SCHEME not in _markdown_to_fragment(markdown)
+
+    # A reply cannot forge one: the scheme is defused like the others.
+    forged = _markdown_to_fragment(f"[copy]({COPY_SCHEME}0.0)", 0)
+    assert f'href="{COPY_SCHEME}' not in forged, forged
+    print("ok  code block copy links")
+
+
 def test_system_qt_theme_guard():
     """The plugin path is only borrowed when the Qt versions agree."""
     import os
@@ -3260,4 +3337,6 @@ if __name__ == "__main__":
     test_title_request()
     test_needs_title()
     test_title_prompt()
+    test_code_block_wrapping()
+    test_code_block_copy_links()
     print("\nall checks passed")
-- 
cgit v1.2.3