Merge pull request #578 from Bortlesboat/feat/555-auto-reply-headers

feat(gmail): auto-populate In-Reply-To/References for reply drafts

gmail/gmail_tools.py

@@ -7,6 +7,7 @@ This module provides MCP tools for interacting with the Gmail API.
 import logging
 import asyncio
 import base64
+import re
 import ssl
 import mimetypes
 from html.parser import HTMLParser
@@ -441,6 +442,91 @@ def _extract_headers(payload: dict, header_names: List[str]) -> Dict[str, str]:
     return headers
 
 
+def _parse_message_id_chain(header_value: Optional[str]) -> List[str]:
+    """Extract Message-IDs from a reply header value."""
+    if not header_value:
+        return []
+
+    message_ids = re.findall(r"<[^>]+>", header_value)
+    if message_ids:
+        return message_ids
+
+    return header_value.split()
+
+
+def _derive_reply_headers(
+    thread_message_ids: List[str],
+    in_reply_to: Optional[str],
+    references: Optional[str],
+) -> tuple[Optional[str], Optional[str]]:
+    """Fill missing reply headers while preserving caller intent."""
+    derived_in_reply_to = in_reply_to
+    derived_references = references
+
+    if not thread_message_ids:
+        return derived_in_reply_to, derived_references
+
+    if not derived_in_reply_to:
+        reference_chain = _parse_message_id_chain(derived_references)
+        derived_in_reply_to = (
+            reference_chain[-1] if reference_chain else thread_message_ids[-1]
+        )
+
+    if not derived_references:
+        if derived_in_reply_to and derived_in_reply_to in thread_message_ids:
+            reply_index = thread_message_ids.index(derived_in_reply_to)
+            derived_references = " ".join(thread_message_ids[: reply_index + 1])
+        elif derived_in_reply_to:
+            derived_references = derived_in_reply_to
+        else:
+            derived_references = " ".join(thread_message_ids)
+
+    return derived_in_reply_to, derived_references
+
+
+async def _fetch_thread_message_ids(service, thread_id: str) -> List[str]:
+    """
+    Fetch Message-ID headers from a Gmail thread for reply threading.
+
+    Args:
+        service: Gmail API service instance
+        thread_id: Gmail thread ID
+
+    Returns:
+        Message-IDs in thread order. Returns an empty list on failure.
+    """
+    try:
+        thread = await asyncio.to_thread(
+            service.users()
+            .threads()
+            .get(
+                userId="me",
+                id=thread_id,
+                format="metadata",
+                metadataHeaders=["Message-ID"],
+            )
+            .execute
+        )
+        messages = thread.get("messages", [])
+        if not messages:
+            return []
+
+        # Collect all Message-IDs in thread order
+        message_ids = []
+        for msg in messages:
+            headers = _extract_headers(msg.get("payload", {}), ["Message-ID"])
+            mid = headers.get("Message-ID")
+            if mid:
+                message_ids.append(mid)
+
+        return message_ids
+    except Exception as e:
+        logger.warning(
+            f"Failed to fetch thread Message-IDs for thread {thread_id}: {e}"
+        )
+        return []
+
+
 def _prepare_gmail_message(
     subject: str,
     body: str,
@@ -1645,6 +1731,14 @@ async def draft_gmail_message(
     else:
         draft_body = _append_signature_to_body(draft_body, body_format, signature_html)
 
+    # Auto-populate In-Reply-To and References when thread_id is provided
+    # but headers are missing, to ensure the draft renders inline in Gmail
+    if thread_id and (not in_reply_to or not references):
+        thread_message_ids = await _fetch_thread_message_ids(service, thread_id)
+        in_reply_to, references = _derive_reply_headers(
+            thread_message_ids, in_reply_to, references
+        )
+
     raw_message, thread_id_final, attached_count = _prepare_gmail_message(
         subject=subject,
         body=draft_body,

tests/gmail/test_draft_gmail_message.py

@@ -0,0 +1,288 @@
+import base64
+import os
+import sys
+from unittest.mock import Mock
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+from core.utils import UserInputError
+from gmail.gmail_tools import draft_gmail_message
+
+
+def _unwrap(tool):
+    """Unwrap FunctionTool + decorators to the original async function."""
+    fn = tool.fn if hasattr(tool, "fn") else tool
+    while hasattr(fn, "__wrapped__"):
+        fn = fn.__wrapped__
+    return fn
+
+
+def _thread_response(*message_ids):
+    return {
+        "messages": [
+            {
+                "payload": {
+                    "headers": [{"name": "Message-ID", "value": message_id}],
+                }
+            }
+            for message_id in message_ids
+        ]
+    }
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_reports_actual_attachment_count(
+    tmp_path, monkeypatch
+):
+    monkeypatch.setenv("ALLOWED_FILE_DIRS", str(tmp_path))
+    attachment_path = tmp_path / "sample.txt"
+    attachment_path.write_text("hello attachment", encoding="utf-8")
+
+    mock_service = Mock()
+    mock_service.users().drafts().create().execute.return_value = {"id": "draft123"}
+
+    result = await _unwrap(draft_gmail_message)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        to="recipient@example.com",
+        subject="Attachment test",
+        body="Please see attached.",
+        attachments=[{"path": str(attachment_path)}],
+        include_signature=False,
+    )
+
+    assert "Draft created with 1 attachment(s)! Draft ID: draft123" in result
+
+    create_kwargs = (
+        mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+    )
+    raw_message = create_kwargs["body"]["message"]["raw"]
+    raw_bytes = base64.urlsafe_b64decode(raw_message)
+
+    assert b"Content-Disposition: attachment;" in raw_bytes
+    assert b"sample.txt" in raw_bytes
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_raises_when_no_attachments_are_added(
+    tmp_path, monkeypatch
+):
+    monkeypatch.setenv("ALLOWED_FILE_DIRS", str(tmp_path))
+    missing_path = tmp_path / "missing.txt"
+
+    mock_service = Mock()
+    mock_service.users().drafts().create().execute.return_value = {"id": "draft123"}
+
+    with pytest.raises(UserInputError, match="No valid attachments were added"):
+        await _unwrap(draft_gmail_message)(
+            service=mock_service,
+            user_google_email="user@example.com",
+            to="recipient@example.com",
+            subject="Attachment test",
+            body="Please see attached.",
+            attachments=[{"path": str(missing_path)}],
+            include_signature=False,
+        )
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_appends_gmail_signature_html():
+    mock_service = Mock()
+    mock_service.users().drafts().create().execute.return_value = {"id": "draft_sig"}
+    mock_service.users().settings().sendAs().list().execute.return_value = {
+        "sendAs": [
+            {
+                "sendAsEmail": "user@example.com",
+                "isPrimary": True,
+                "signature": "<div>Best,<br>Alice</div>",
+            }
+        ]
+    }
+
+    result = await _unwrap(draft_gmail_message)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        to="recipient@example.com",
+        subject="Signature test",
+        body="<p>Hello</p>",
+        body_format="html",
+        include_signature=True,
+    )
+
+    assert "Draft created! Draft ID: draft_sig" in result
+
+    create_kwargs = (
+        mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+    )
+    raw_message = create_kwargs["body"]["message"]["raw"]
+    raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+    assert "<p>Hello</p>" in raw_text
+    assert "Best,<br>Alice" in raw_text
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_autofills_reply_headers_from_thread():
+    mock_service = Mock()
+    mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+    mock_service.users().threads().get().execute.return_value = _thread_response(
+        "<msg1@example.com>",
+        "<msg2@example.com>",
+        "<msg3@example.com>",
+    )
+
+    result = await _unwrap(draft_gmail_message)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        to="recipient@example.com",
+        subject="Meeting tomorrow",
+        body="Thanks for the update.",
+        thread_id="thread123",
+        include_signature=False,
+    )
+
+    # Verify threads().get() was called with correct parameters
+    thread_get_kwargs = (
+        mock_service.users.return_value.threads.return_value.get.call_args.kwargs
+    )
+    assert thread_get_kwargs["userId"] == "me"
+    assert thread_get_kwargs["id"] == "thread123"
+    assert thread_get_kwargs["format"] == "metadata"
+    assert "Message-ID" in thread_get_kwargs["metadataHeaders"]
+
+    assert "Draft created! Draft ID: draft_reply" in result
+
+    create_kwargs = (
+        mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+    )
+    raw_message = create_kwargs["body"]["message"]["raw"]
+    raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+    assert "In-Reply-To: <msg3@example.com>" in raw_text
+    assert (
+        "References: <msg1@example.com> <msg2@example.com> <msg3@example.com>"
+        in raw_text
+    )
+    assert create_kwargs["body"]["message"]["threadId"] == "thread123"
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_uses_explicit_in_reply_to_when_filling_references():
+    mock_service = Mock()
+    mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+    mock_service.users().threads().get().execute.return_value = _thread_response(
+        "<msg1@example.com>",
+        "<msg2@example.com>",
+        "<msg3@example.com>",
+    )
+
+    await _unwrap(draft_gmail_message)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        to="recipient@example.com",
+        subject="Meeting tomorrow",
+        body="Replying to an earlier message.",
+        thread_id="thread123",
+        in_reply_to="<msg2@example.com>",
+        include_signature=False,
+    )
+
+    create_kwargs = (
+        mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+    )
+    raw_message = create_kwargs["body"]["message"]["raw"]
+    raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+    assert "In-Reply-To: <msg2@example.com>" in raw_text
+    assert "References: <msg1@example.com> <msg2@example.com>" in raw_text
+    assert "<msg3@example.com>" not in raw_text
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_uses_explicit_references_when_filling_in_reply_to():
+    mock_service = Mock()
+    mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+    mock_service.users().threads().get().execute.return_value = _thread_response(
+        "<msg1@example.com>",
+        "<msg2@example.com>",
+        "<msg3@example.com>",
+    )
+
+    await _unwrap(draft_gmail_message)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        to="recipient@example.com",
+        subject="Meeting tomorrow",
+        body="Replying to an earlier message.",
+        thread_id="thread123",
+        references="<msg1@example.com> <msg2@example.com>",
+        include_signature=False,
+    )
+
+    create_kwargs = (
+        mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+    )
+    raw_message = create_kwargs["body"]["message"]["raw"]
+    raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+    assert "In-Reply-To: <msg2@example.com>" in raw_text
+    assert "References: <msg1@example.com> <msg2@example.com>" in raw_text
+    assert "<msg3@example.com>" not in raw_text
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_gracefully_degrades_when_thread_fetch_fails():
+    mock_service = Mock()
+    mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+    mock_service.users().threads().get().execute.side_effect = RuntimeError("boom")
+
+    result = await _unwrap(draft_gmail_message)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        to="recipient@example.com",
+        subject="Meeting tomorrow",
+        body="Thanks for the update.",
+        thread_id="thread123",
+        include_signature=False,
+    )
+
+    assert "Draft created! Draft ID: draft_reply" in result
+
+    create_kwargs = (
+        mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+    )
+    raw_message = create_kwargs["body"]["message"]["raw"]
+    raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+    assert "In-Reply-To:" not in raw_text
+    assert "References:" not in raw_text
+
+
+@pytest.mark.asyncio
+async def test_draft_gmail_message_gracefully_degrades_when_thread_has_no_messages():
+    mock_service = Mock()
+    mock_service.users().drafts().create().execute.return_value = {"id": "draft_reply"}
+    mock_service.users().threads().get().execute.return_value = {"messages": []}
+
+    result = await _unwrap(draft_gmail_message)(
+        service=mock_service,
+        user_google_email="user@example.com",
+        to="recipient@example.com",
+        subject="Meeting tomorrow",
+        body="Thanks for the update.",
+        thread_id="thread123",
+        include_signature=False,
+    )
+
+    assert "Draft created! Draft ID: draft_reply" in result
+
+    create_kwargs = (
+        mock_service.users.return_value.drafts.return_value.create.call_args.kwargs
+    )
+    raw_message = create_kwargs["body"]["message"]["raw"]
+    raw_text = base64.urlsafe_b64decode(raw_message).decode("utf-8", errors="ignore")
+
+    assert "In-Reply-To:" not in raw_text
+    assert "References:" not in raw_text