Skip to content

Handle HTTP-date Retry-After at Graph call sites #95

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gaoflow wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Handle HTTP-date Retry-After at Graph call sites #95

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/entrabot/graph_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

import requests

from entrabot.tools.rate_limit import parse_retry_after
from entrabot.url_safety import _is_graph_url

GRAPH_BETA = "https://graph.microsoft.com/beta"
Expand Down Expand Up @@ -84,7 +85,7 @@ def graph_request(
resp = requests.request(method, url, headers=headers, json=json_body, timeout=timeout)

if retry and resp.status_code in _RETRYABLE:
wait = int(resp.headers.get("Retry-After", "10"))
wait = parse_retry_after(resp.headers.get("Retry-After"), default=10)
print(f" Graph API returned {resp.status_code}; retrying in {wait}s…")
time.sleep(wait)
resp = requests.request(method, url, headers=headers, json=json_body, timeout=timeout)
Expand Down
49 changes: 19 additions & 30 deletions src/entrabot/tools/teams.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
from entrabot.graph_helpers import odata_escape
from entrabot.platform import get_credential_store
from entrabot.tools.audit import log_event
from entrabot.tools.rate_limit import RetryOn429Transport
from entrabot.tools.rate_limit import RetryOn429Transport, parse_retry_after
from entrabot.url_safety import _is_graph_url

logger = logging.getLogger("entrabot.tools.teams")
Expand Down Expand Up @@ -336,9 +336,7 @@ async def create_one_on_one_chat(
agent_member: dict = {
"@odata.type": "#microsoft.graph.aadUserConversationMember",
"roles": ["owner"],
"user@odata.bind": (
f"https://graph.microsoft.com/v1.0/users('{_agent_bind_value}')"
),
"user@odata.bind": (f"https://graph.microsoft.com/v1.0/users('{_agent_bind_value}')"),
}

payload = {
Expand Down Expand Up @@ -371,7 +369,7 @@ async def create_one_on_one_chat(
logger.error("400 creating 1:1 chat: %s", error_msg)
raise ValueError(f"Graph API rejected chat creation: {error_msg}")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
resp.raise_for_status()

Expand Down Expand Up @@ -499,7 +497,7 @@ async def create_or_find_chat(
"their actual UPN. Check ENTRABOT_HUMAN_USER_MAILS in .env."
)
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
resp.raise_for_status()

Expand Down Expand Up @@ -635,9 +633,7 @@ async def add_member(
(
s
for s in sponsors
if any(
e.lower() == requester_lower for e in s.email_identifiers()
)
if any(e.lower() == requester_lower for e in s.email_identifiers())
),
None,
)
Expand Down Expand Up @@ -711,13 +707,8 @@ async def add_member(
# Gate 2 (defense-in-depth): existing Graph membership check.
members = await _fetch_chat_members_for_gate(chat_id)
matched_user_id = matched_sponsor.user_id.lower()
if not any(
(m.get("user_id") or "").strip().lower() == matched_user_id
for m in members
):
err = RequesterNotInChatError(
requester=requester_email, chat_id=chat_id
)
if not any((m.get("user_id") or "").strip().lower() == matched_user_id for m in members):
err = RequesterNotInChatError(requester=requester_email, chat_id=chat_id)
log_event(
action="teams.add_member",
resource=audit_resource,
Expand Down Expand Up @@ -771,11 +762,9 @@ async def add_member(
error_msg = resp.text or "Not found"
raise ChatNotFound(f"Could not add member: {error_msg}")
if resp.status_code == 401:
raise TokenExpiredError(
"Agent User token expired — re-acquire via three-hop flow"
)
raise TokenExpiredError("Agent User token expired — re-acquire via three-hop flow")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
resp.raise_for_status()

Expand Down Expand Up @@ -842,7 +831,7 @@ async def list_members(
if resp.status_code == 401:
raise TokenExpiredError("Agent User token expired — re-acquire via three-hop flow")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
resp.raise_for_status()

Expand Down Expand Up @@ -929,7 +918,7 @@ async def send(
if resp.status_code == 401:
raise TokenExpiredError("Token expired — re-acquire")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
if resp.status_code == 404:
raise ChatNotFound(f"Chat {chat_id} not found")
Expand Down Expand Up @@ -985,7 +974,7 @@ async def post_thinking_placeholder(
if resp.status_code == 404:
raise ChatNotFound(f"Chat {chat_id} not found")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
resp.raise_for_status()

Expand Down Expand Up @@ -1039,7 +1028,7 @@ async def update_placeholder(
if resp.status_code == 401:
raise TokenExpiredError("Token expired — re-acquire")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
if 200 <= resp.status_code < 300:
return {"message_id": placeholder_id, "mode": "edit"}
Expand Down Expand Up @@ -1098,7 +1087,7 @@ async def resolve_placeholder(
if resp.status_code == 401:
raise TokenExpiredError("Token expired — re-acquire")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
if 200 <= resp.status_code < 300:
return {"message_id": placeholder_id, "mode": "edit"}
Expand All @@ -1121,7 +1110,7 @@ async def resolve_placeholder(
if sd_resp.status_code == 401:
raise TokenExpiredError("Token expired — re-acquire")
if sd_resp.status_code == 429:
retry_after = int(sd_resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(sd_resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
delete_ok = 200 <= sd_resp.status_code < 300
if not delete_ok:
Expand Down Expand Up @@ -1198,7 +1187,7 @@ async def delete_chat_message(
if resp.status_code == 401:
raise TokenExpiredError("Token expired — re-acquire")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
if 200 <= resp.status_code < 300:
return True
Expand Down Expand Up @@ -1234,7 +1223,7 @@ async def fetch_hosted_image(*, token: str, url: str) -> bytes | None:
if resp.status_code == 404:
return None
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
resp.raise_for_status()
return resp.content
Expand Down Expand Up @@ -1324,7 +1313,7 @@ async def fetch_message(
if resp.status_code == 401:
raise TokenExpiredError("Token expired — re-acquire")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
if not (200 <= resp.status_code < 300):
return None
Expand Down Expand Up @@ -1366,7 +1355,7 @@ async def read(
if resp.status_code == 404:
raise ChatNotFound(f"Chat {chat_id} not found")
if resp.status_code == 429:
retry_after = int(resp.headers.get("Retry-After", "60"))
retry_after = parse_retry_after(resp.headers.get("Retry-After"), default=60)
raise RateLimitError(retry_after)
resp.raise_for_status()

Expand Down
23 changes: 20 additions & 3 deletions tests/test_graph_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,22 @@ def test_retry_on_429(self):
assert mock_requests.request.call_count == 2
mock_time.sleep.assert_called_once_with(1)

def test_retry_on_429_with_http_date_retry_after(self):
from entrabot.graph_helpers import graph_request

throttled = _resp(429, headers={"Retry-After": "Wed, 21 Oct 2015 07:28:00 GMT"})
ok = _resp(200, {"value": []})
with (
patch("entrabot.graph_helpers.requests") as mock_requests,
patch("entrabot.graph_helpers.time") as mock_time,
):
mock_requests.request.side_effect = [throttled, ok]
result = graph_request("GET", "/users", "tok")

assert result.status_code == 200
assert mock_requests.request.call_count == 2
mock_time.sleep.assert_called_once_with(0)

def test_retry_on_503(self):
from entrabot.graph_helpers import graph_request

Expand Down Expand Up @@ -260,9 +276,10 @@ def test_accepts_commercial_graph_next_link_and_fetches_second_page(self):

assert result == [{"id": "a"}, {"id": "b"}]
assert mock_requests.request.call_args_list[1].args == ("GET", next_link)
assert mock_requests.request.call_args_list[1].kwargs["headers"][
"Authorization"
] == "Bearer tok"
assert (
mock_requests.request.call_args_list[1].kwargs["headers"]["Authorization"]
== "Bearer tok"
)

def test_error_raises(self):
from entrabot.graph_helpers import graph_collection_values
Expand Down
25 changes: 22 additions & 3 deletions tests/tools/test_teams.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
AgentIDNotAvailable,
ChatNotFound,
MessageTooLong,
RateLimitError,
TeamsNotLicensed,
TokenExchangeError,
TokenExpiredError,
Expand Down Expand Up @@ -765,9 +766,7 @@ def _sponsor_gate_patches(*, chat_id: str = "19:chat-id@thread.v2"):
patch(
"entrabot.tools.teams._fetch_chat_members_for_gate",
new=AsyncMock(
return_value=[
{"user_id": "sponsor-uid", "email": "sponsor@contoso.com"}
]
return_value=[{"user_id": "sponsor-uid", "email": "sponsor@contoso.com"}]
),
),
]
Expand Down Expand Up @@ -1440,6 +1439,26 @@ async def test_not_licensed(self) -> None:
agent_user_id="agent-oid-123",
)

@respx.mock
@pytest.mark.asyncio
async def test_429_http_date_retry_after_raises_rate_limit(self) -> None:
from entrabot.tools.teams import create_one_on_one_chat

respx.post(f"{GRAPH_BASE}/chats").mock(
return_value=httpx.Response(
429,
headers={"Retry-After": "Wed, 21 Oct 2015 07:28:00 GMT"},
)
)
with pytest.raises(RateLimitError) as exc_info:
await create_one_on_one_chat(
token="tok",
target_email="someone@example.com",
agent_user_id="agent-oid-123",
)

assert exc_info.value.retry_after == 0


# ---------------------------------------------------------------------------
# fetch_hosted_image
Expand Down