human hand off

backend/app/services/mime_body_extractor.py

@@ -2,6 +2,10 @@ from dataclasses import dataclass, field
 from email import policy
 from email.message import EmailMessage, Message
 from email.parser import BytesParser
+from html.parser import HTMLParser
+import re
+
+from app.services.reply_extractor import extract_new_reply_content
 
 
 @dataclass
@@ -104,3 +108,109 @@ def extract_mirror_content(raw: bytes) -> MirroredMimeContent:
     collector = MirroredMimeContent()
     _walk_parts(message, collector)
     return collector
+
+
+class _NewReplyHTMLParser(HTMLParser):
+    QUOTE_CLASSES = {
+        "gmail_attr",
+        "gmail_quote",
+        "protonmail_quote",
+        "yahoo_quoted",
+        "moz-cite-prefix",
+    }
+    QUOTE_IDS = {"divRplyFwdMsg", "appendonsend"}
+    VOID_TAGS = {"br", "hr", "img", "meta", "link", "input"}
+
+    def __init__(self) -> None:
+        super().__init__(convert_charrefs=False)
+        self.output: list[str] = []
+        self.open_tags: list[str] = []
+        self.stopped = False
+
+    def _is_quote_start(self, tag: str, attrs: list[tuple[str, str | None]]) -> bool:
+        values = {key.lower(): (value or "") for key, value in attrs}
+        classes = set(values.get("class", "").split())
+        return (
+            tag.lower() == "blockquote"
+            or bool(classes.intersection(self.QUOTE_CLASSES))
+            or values.get("id") in self.QUOTE_IDS
+        )
+
+    def handle_starttag(self, tag: str, attrs) -> None:
+        if self.stopped:
+            return
+        if self._is_quote_start(tag, attrs):
+            self.stopped = True
+            return
+        self.output.append(self.get_starttag_text())
+        if tag.lower() not in self.VOID_TAGS:
+            self.open_tags.append(tag)
+
+    def handle_startendtag(self, tag: str, attrs) -> None:
+        if not self.stopped:
+            self.output.append(self.get_starttag_text())
+
+    def handle_endtag(self, tag: str) -> None:
+        if self.stopped:
+            return
+        self.output.append(f"</{tag}>")
+        if tag in self.open_tags:
+            reverse_index = self.open_tags[::-1].index(tag)
+            del self.open_tags[len(self.open_tags) - reverse_index - 1]
+
+    def handle_data(self, data: str) -> None:
+        if self.stopped:
+            return
+        if re.search(r"(?:^|\n)\s*On .+? wrote:\s*$", data, re.IGNORECASE):
+            self.stopped = True
+            return
+        self.output.append(data)
+
+    def handle_entityref(self, name: str) -> None:
+        if not self.stopped:
+            self.output.append(f"&{name};")
+
+    def handle_charref(self, name: str) -> None:
+        if not self.stopped:
+            self.output.append(f"&#{name};")
+
+    def html(self) -> str:
+        value = "".join(self.output).strip()
+        if not value:
+            return ""
+        return value + "".join(f"</{tag}>" for tag in reversed(self.open_tags))
+
+
+def _new_reply_html(html: str | None) -> str | None:
+    if not html:
+        return None
+    parser = _NewReplyHTMLParser()
+    parser.feed(html)
+    parser.close()
+    value = parser.html().strip()
+    return value or None
+
+
+def extract_new_reply_mime_content(raw: bytes) -> MirroredMimeContent:
+    """Extract only newly authored MIME content while retaining rich formatting."""
+    full = extract_mirror_content(raw)
+    extracted = extract_new_reply_content(
+        body_text=full.body_text,
+        body_html=full.body_html,
+    )
+    new_html = _new_reply_html(full.body_html)
+    referenced_cids = {
+        match.casefold()
+        for match in re.findall(r"cid:([^\"'\s>)]+)", new_html or "", re.IGNORECASE)
+    }
+    inline_parts = [
+        part
+        for part in full.inline_parts
+        if part.content_id and part.content_id.casefold() in referenced_cids
+    ]
+    return MirroredMimeContent(
+        body_text=extracted["new_reply_content"].strip(),
+        body_html=new_html,
+        inline_parts=inline_parts,
+        file_attachments=full.file_attachments,
+    )