diff --git a/.fernignore b/.fernignore index 950b552..206f5a6 100644 --- a/.fernignore +++ b/.fernignore @@ -2,3 +2,8 @@ src/agentmail/core/query_encoder.py tests/utils/test_query_encoding.py + +# Frozen to preserve the Idempotency-Key auto-mint patch (see idempotency.py). +src/agentmail/core/idempotency.py +src/agentmail/core/http_client.py +tests/utils/test_idempotency.py diff --git a/src/agentmail/core/http_client.py b/src/agentmail/core/http_client.py index ee93758..b20d74e 100644 --- a/src/agentmail/core/http_client.py +++ b/src/agentmail/core/http_client.py @@ -11,6 +11,7 @@ import httpx from .file import File, convert_file_dict_to_httpx_tuples from .force_multipart import FORCE_MULTIPART +from .idempotency import maybe_mint_idempotency_key from .jsonable_encoder import jsonable_encoder from .logging import LogConfig, Logger, create_logger from .query_encoder import encode_query @@ -304,6 +305,20 @@ def request( omit: typing.Optional[typing.Any] = None, force_multipart: typing.Optional[bool] = None, ) -> httpx.Response: + # Auto-mint an Idempotency-Key for send endpoints. This runs ONCE per + # logical call (retries == 0, before the retry recursion below) so every + # internal retry reuses the same key. The minted key is stored on the + # `headers` local, which is re-passed as headers=headers in the recursion. + if retries == 0: + _minted_headers = maybe_mint_idempotency_key( + method=method, + path=path, + headers=headers, + additional_headers=(request_options.get("additional_headers") if request_options is not None else None), + ) + if _minted_headers is not None: + headers = _minted_headers + base_url = self.get_base_url(base_url) timeout = ( request_options.get("timeout_in_seconds") @@ -564,6 +579,20 @@ async def request( omit: typing.Optional[typing.Any] = None, force_multipart: typing.Optional[bool] = None, ) -> httpx.Response: + # Auto-mint an Idempotency-Key for send endpoints. This runs ONCE per + # logical call (retries == 0, before the retry recursion below) so every + # internal retry reuses the same key. The minted key is stored on the + # `headers` local, which is re-passed as headers=headers in the recursion. + if retries == 0: + _minted_headers = maybe_mint_idempotency_key( + method=method, + path=path, + headers=headers, + additional_headers=(request_options.get("additional_headers") if request_options is not None else None), + ) + if _minted_headers is not None: + headers = _minted_headers + base_url = self.get_base_url(base_url) timeout = ( request_options.get("timeout_in_seconds") diff --git a/src/agentmail/core/idempotency.py b/src/agentmail/core/idempotency.py new file mode 100644 index 0000000..2605997 --- /dev/null +++ b/src/agentmail/core/idempotency.py @@ -0,0 +1,74 @@ +# This file is a hand-written AgentMail patch (NOT Fern-generated). It is pinned +# in .fernignore so Fern regeneration cannot clobber it. +# +# Purpose: auto-mint an `Idempotency-Key` header for send endpoints so that the +# HttpClient's internal auto-retries (429/5xx/timeouts) all carry the SAME key. +# The AgentMail API makes SENDS idempotent via this header: a retry with the same +# key returns the original message (no duplicate email); reuse with a different +# payload is a 409. Minting once per logical call (before the retry loop) keeps +# every attempt idempotent. + +import re +import typing +import uuid + +IDEMPOTENCY_KEY_HEADER = "Idempotency-Key" + +# The 5 POST send endpoints, matched on the path suffix (after the /v0 prefix and +# ignoring any query string). {segment} is a single non-slash path segment. +# +# Anchored to the END of the path ($) so that ONLY the exact final segments +# /send, /reply, /reply-all, /forward match. This deliberately does NOT match the +# create-draft endpoints (/messages/{seg}/draft-reply, /draft-reply-all, +# /draft-forward), which end in different segments. +_SEND_PATH_REGEX = re.compile( + r"(?:^|/)(?:" + r"messages/send" + r"|messages/[^/]+/reply" + r"|messages/[^/]+/reply-all" + r"|messages/[^/]+/forward" + r"|drafts/[^/]+/send" + r")$" +) + + +def _is_send_path(path: typing.Optional[str]) -> bool: + if not path: + return False + # Ignore any query string / fragment; strip a trailing slash before anchoring. + clean = path.split("?", 1)[0].split("#", 1)[0].rstrip("/") + return _SEND_PATH_REGEX.search(clean) is not None + + +def _has_idempotency_key(headers: typing.Optional[typing.Mapping[str, typing.Any]]) -> bool: + if not headers: + return False + return any(k.lower() == IDEMPOTENCY_KEY_HEADER.lower() for k in headers.keys()) + + +def maybe_mint_idempotency_key( + *, + method: str, + path: typing.Optional[str], + headers: typing.Optional[typing.Dict[str, typing.Any]], + additional_headers: typing.Optional[typing.Mapping[str, typing.Any]], +) -> typing.Optional[typing.Dict[str, typing.Any]]: + """Return a headers dict carrying an auto-minted Idempotency-Key, or None. + + Mints (UUID4) only when: the method is POST, the path is one of the send + endpoints, and no caller-supplied `Idempotency-Key` header is already present + (case-insensitive) in either `headers` or the request_options + `additional_headers`. Returns a NEW dict (never mutates the caller's) that the + request method should thread through its retry recursion so all attempts share + the same key. Returns None when nothing should change. + """ + if method.upper() != "POST": + return None + if not _is_send_path(path): + return None + if _has_idempotency_key(headers) or _has_idempotency_key(additional_headers): + return None + + minted = dict(headers) if headers else {} + minted[IDEMPOTENCY_KEY_HEADER] = str(uuid.uuid4()) + return minted diff --git a/tests/utils/test_idempotency.py b/tests/utils/test_idempotency.py new file mode 100644 index 0000000..5ade74d --- /dev/null +++ b/tests/utils/test_idempotency.py @@ -0,0 +1,172 @@ +# Tests for the hand-written Idempotency-Key auto-mint patch. Pinned in +# .fernignore so Fern regeneration cannot clobber it. + +import uuid +from typing import List + +import httpx +import pytest + +from agentmail.core.http_client import HttpClient +from agentmail.core.request_options import RequestOptions + + +def _is_valid_uuid4(value: str) -> bool: + try: + parsed = uuid.UUID(value) + except (ValueError, TypeError, AttributeError): + return False + return parsed.version == 4 and str(parsed) == value.lower() + + +def _make_client(handler: httpx.MockTransport) -> HttpClient: + return HttpClient( + httpx_client=httpx.Client(transport=handler), + base_timeout=lambda: None, + base_headers=lambda: {}, + base_url=lambda: "https://api.agentmail.to/v0", + ) + + +def _capture_transport(captured: List[httpx.Request], status_codes: List[int]) -> httpx.MockTransport: + """MockTransport that records each request and returns the next status code.""" + calls = {"i": 0} + + def handle(request: httpx.Request) -> httpx.Response: + captured.append(request) + idx = min(calls["i"], len(status_codes) - 1) + calls["i"] += 1 + return httpx.Response(status_codes[idx], json={}) + + return httpx.MockTransport(handle) + + +def test_send_path_mints_uuid4_idempotency_key() -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + client.request(path="messages/send", method="POST", json={"to": "a@b.com"}, request_options=None) + + assert len(captured) == 1 + key = captured[0].headers.get("Idempotency-Key") + assert key is not None + assert _is_valid_uuid4(key) + + +def test_caller_supplied_key_is_preserved() -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + options: RequestOptions = {"additional_headers": {"Idempotency-Key": "caller-key-123"}} + client.request(path="messages/send", method="POST", json={"to": "a@b.com"}, request_options=options) + + assert len(captured) == 1 + assert captured[0].headers.get("Idempotency-Key") == "caller-key-123" + + +def test_retry_reuses_identical_key() -> None: + captured: List[httpx.Request] = [] + # 503 once, then 200 -> the client should retry and both attempts share the key. + client = _make_client(_capture_transport(captured, [503, 200])) + + client.request( + path="messages/send", + method="POST", + json={"to": "a@b.com"}, + request_options={"max_retries": 2}, + ) + + assert len(captured) == 2 + keys = [r.headers.get("Idempotency-Key") for r in captured] + assert keys[0] is not None + assert _is_valid_uuid4(keys[0]) + assert keys[0] == keys[1] + + +def test_non_send_post_is_not_injected() -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + client.request(path="inboxes", method="POST", json={"username": "x"}, request_options=None) + + assert len(captured) == 1 + assert "Idempotency-Key" not in captured[0].headers + + +def test_two_calls_mint_different_keys() -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + client.request(path="messages/send", method="POST", json={"to": "a@b.com"}, request_options=None) + client.request(path="messages/send", method="POST", json={"to": "c@d.com"}, request_options=None) + + assert len(captured) == 2 + key1 = captured[0].headers.get("Idempotency-Key") + key2 = captured[1].headers.get("Idempotency-Key") + assert key1 is not None and key2 is not None + assert key1 != key2 + + +def test_create_draft_path_is_not_injected() -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + # /messages/{id}/draft-reply is a create-draft endpoint, NOT a send. + client.request(path="messages/msg_123/draft-reply", method="POST", json={}, request_options=None) + + assert len(captured) == 1 + assert "Idempotency-Key" not in captured[0].headers + + +@pytest.mark.parametrize( + "path", + [ + "messages/send", + "messages/msg_1/reply", + "messages/msg_1/reply-all", + "messages/msg_1/forward", + "drafts/dft_1/send", + ], +) +def test_all_send_paths_are_injected(path: str) -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + client.request(path=path, method="POST", json={}, request_options=None) + + assert "Idempotency-Key" in captured[0].headers, f"expected injection for {path}" + + +@pytest.mark.parametrize( + "path", + [ + "messages/msg_1/draft-reply", + "messages/msg_1/draft-reply-all", + "messages/msg_1/draft-forward", + ], +) +def test_create_draft_paths_are_not_injected(path: str) -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + client.request(path=path, method="POST", json={}, request_options=None) + + assert "Idempotency-Key" not in captured[0].headers, f"unexpected injection for {path}" + + +def test_send_path_with_query_string_is_injected() -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + client.request(path="messages/send?foo=bar", method="POST", json={}, request_options=None) + + assert "Idempotency-Key" in captured[0].headers + + +def test_get_on_send_path_is_not_injected() -> None: + captured: List[httpx.Request] = [] + client = _make_client(_capture_transport(captured, [200])) + + client.request(path="messages/send", method="GET", request_options=None) + + assert "Idempotency-Key" not in captured[0].headers