mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-05-10 03:22:05 +00:00
fix(msgraph): stream download_to_file body instead of buffering
The prior implementation routed download_to_file through the shared _request() path, which uses httpx.AsyncClient.request() inside a context manager that closes before aiter_bytes() iterates. The body was read into memory first and the chunked write loop replayed it from buffer. On small test payloads this was invisible; on real Teams meeting recordings (hundreds of MB) it would force the full artifact into RAM per download. Rewrites download_to_file to open its own AsyncClient and use client.stream(), keeping the context open across the aiter_bytes iteration so the body is actually streamed chunk-by-chunk to disk. Retry/token-refresh/Retry-After semantics are preserved by handling them inline on the stream path. Partial .part files are cleaned up on transport errors and on exhausted retries. Adds three tests: large-payload streaming verifies the chunk loop runs multiple times (discriminator: 512 KiB at chunk_size=65536 yields 8 chunks under streaming, 1 under buffering), transient-5xx retry recovers after a single retry, and exhausted-retry cleans up the partial file.
This commit is contained in:
parent
b878f89f66
commit
45d860d424
2 changed files with 197 additions and 11 deletions
|
|
@ -160,20 +160,101 @@ class MicrosoftGraphClient:
|
|||
headers: dict[str, str] | None = None,
|
||||
chunk_size: int = 65536,
|
||||
) -> dict[str, Any]:
|
||||
response = await self._request("GET", path, headers=headers)
|
||||
"""Download a Graph resource to disk, streaming the response body.
|
||||
|
||||
The body is written chunk-by-chunk via ``response.aiter_bytes`` with
|
||||
the ``httpx.AsyncClient`` kept open for the duration of the iteration,
|
||||
so recordings and other large artifacts do not need to fit in memory.
|
||||
"""
|
||||
url = self._resolve_url(path)
|
||||
target = Path(destination)
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_target = target.with_suffix(target.suffix + ".part")
|
||||
with tmp_target.open("wb") as handle:
|
||||
async for chunk in response.aiter_bytes(chunk_size=chunk_size):
|
||||
if chunk:
|
||||
handle.write(chunk)
|
||||
os.replace(tmp_target, target)
|
||||
return {
|
||||
"path": str(target),
|
||||
"size_bytes": target.stat().st_size,
|
||||
"content_type": response.headers.get("content-type"),
|
||||
}
|
||||
|
||||
attempt = 0
|
||||
last_error: Exception | None = None
|
||||
|
||||
while attempt <= self.max_retries:
|
||||
token = await self.token_provider.get_access_token(
|
||||
force_refresh=attempt > 0 and self._should_refresh_token(last_error)
|
||||
)
|
||||
request_headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": self.user_agent,
|
||||
}
|
||||
if headers:
|
||||
request_headers.update(headers)
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self.timeout),
|
||||
transport=self._transport,
|
||||
) as client:
|
||||
async with client.stream(
|
||||
"GET",
|
||||
url,
|
||||
headers=request_headers,
|
||||
) as response:
|
||||
if response.status_code >= 400:
|
||||
# Materialize error body so we can surface a meaningful
|
||||
# message; error bodies are small.
|
||||
await response.aread()
|
||||
api_error = self._build_api_error("GET", url, response)
|
||||
last_error = api_error
|
||||
|
||||
if (
|
||||
response.status_code == 401
|
||||
and attempt < self.max_retries
|
||||
):
|
||||
self.token_provider.clear_cache()
|
||||
await self._sleep(
|
||||
self._retry_delay(response, attempt)
|
||||
)
|
||||
attempt += 1
|
||||
continue
|
||||
|
||||
if (
|
||||
self._should_retry(response)
|
||||
and attempt < self.max_retries
|
||||
):
|
||||
await self._sleep(
|
||||
self._retry_delay(response, attempt)
|
||||
)
|
||||
attempt += 1
|
||||
continue
|
||||
|
||||
raise api_error
|
||||
|
||||
content_type = response.headers.get("content-type")
|
||||
with tmp_target.open("wb") as handle:
|
||||
async for chunk in response.aiter_bytes(
|
||||
chunk_size=chunk_size
|
||||
):
|
||||
if chunk:
|
||||
handle.write(chunk)
|
||||
except httpx.HTTPError as exc:
|
||||
last_error = exc
|
||||
tmp_target.unlink(missing_ok=True)
|
||||
if attempt >= self.max_retries:
|
||||
raise MicrosoftGraphClientError(
|
||||
f"Microsoft Graph download failed for GET {url}: {exc}"
|
||||
) from exc
|
||||
await self._sleep(self._retry_delay(None, attempt))
|
||||
attempt += 1
|
||||
continue
|
||||
|
||||
os.replace(tmp_target, target)
|
||||
return {
|
||||
"path": str(target),
|
||||
"size_bytes": target.stat().st_size,
|
||||
"content_type": content_type,
|
||||
}
|
||||
|
||||
tmp_target.unlink(missing_ok=True)
|
||||
raise MicrosoftGraphClientError(
|
||||
f"Microsoft Graph download exhausted retries for GET {url}."
|
||||
)
|
||||
|
||||
async def _request(
|
||||
self,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue