diff --git a/gmail/gmail_tools.py b/gmail/gmail_tools.py index 1f4f7c7..75afc41 100644 --- a/gmail/gmail_tools.py +++ b/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, diff --git a/tests/gmail/test_draft_gmail_message.py b/tests/gmail/test_draft_gmail_message.py new file mode 100644 index 0000000..fae369a --- /dev/null +++ b/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": "
Best,
Alice
", + } + ] + } + + result = await _unwrap(draft_gmail_message)( + service=mock_service, + user_google_email="user@example.com", + to="recipient@example.com", + subject="Signature test", + body="

Hello

", + 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 "

Hello

" in raw_text + assert "Best,
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( + "", + "", + "", + ) + + 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: " in raw_text + assert ( + "References: " + 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( + "", + "", + "", + ) + + 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="", + 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: " in raw_text + assert "References: " in raw_text + assert "" 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( + "", + "", + "", + ) + + 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=" ", + 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: " in raw_text + assert "References: " in raw_text + assert "" 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