mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-08 13:12:08 +00:00
fix(discord): bound standalone response reads
This commit is contained in:
parent
87be36c240
commit
e0bca1cbe2
2 changed files with 208 additions and 9 deletions
|
|
@ -7433,6 +7433,8 @@ if DISCORD_AVAILABLE:
|
|||
# same channel on every send when the directory cache has no entry (e.g. fresh
|
||||
# install, or channel created after the last directory build).
|
||||
_DISCORD_CHANNEL_TYPE_PROBE_CACHE: Dict[str, bool] = {}
|
||||
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES = 1 * 1024 * 1024
|
||||
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES = 8 * 1024
|
||||
|
||||
|
||||
def _remember_channel_is_forum(chat_id: str, is_forum: bool) -> None:
|
||||
|
|
@ -7469,6 +7471,91 @@ def _standalone_sanitize_error(text) -> str:
|
|||
)
|
||||
|
||||
|
||||
def _standalone_is_mock_object(value: Any) -> bool:
|
||||
return value.__class__.__module__.startswith("unittest.mock")
|
||||
|
||||
|
||||
def _standalone_close_response(resp: Any) -> None:
|
||||
close = getattr(resp, "close", None)
|
||||
if callable(close):
|
||||
close()
|
||||
return
|
||||
release = getattr(resp, "release", None)
|
||||
if callable(release):
|
||||
release()
|
||||
|
||||
|
||||
async def _standalone_read_response_bytes_limited(
|
||||
resp: Any,
|
||||
limit_bytes: int,
|
||||
) -> Tuple[Optional[bytes], bool]:
|
||||
content = getattr(resp, "content", None)
|
||||
if content is None or _standalone_is_mock_object(content):
|
||||
return None, False
|
||||
|
||||
read = getattr(content, "read", None)
|
||||
if callable(read):
|
||||
chunks: list[bytes] = []
|
||||
total = 0
|
||||
while total <= limit_bytes:
|
||||
chunk = await read(limit_bytes + 1 - total)
|
||||
if not chunk:
|
||||
break
|
||||
if isinstance(chunk, str):
|
||||
chunk = chunk.encode("utf-8", "replace")
|
||||
total += len(chunk)
|
||||
chunks.append(chunk)
|
||||
if total > limit_bytes:
|
||||
_standalone_close_response(resp)
|
||||
return b"".join(chunks)[:limit_bytes], True
|
||||
return b"".join(chunks), False
|
||||
|
||||
iter_chunked = getattr(content, "iter_chunked", None)
|
||||
if callable(iter_chunked):
|
||||
chunks: list[bytes] = []
|
||||
total = 0
|
||||
async for chunk in iter_chunked(limit_bytes + 1):
|
||||
if isinstance(chunk, str):
|
||||
chunk = chunk.encode("utf-8", "replace")
|
||||
total += len(chunk)
|
||||
chunks.append(chunk)
|
||||
if total > limit_bytes:
|
||||
_standalone_close_response(resp)
|
||||
return b"".join(chunks)[:limit_bytes], True
|
||||
return b"".join(chunks), False
|
||||
|
||||
return None, False
|
||||
|
||||
|
||||
def _standalone_response_encoding(resp: Any) -> str:
|
||||
get_encoding = getattr(resp, "get_encoding", None)
|
||||
if callable(get_encoding):
|
||||
try:
|
||||
return get_encoding() or "utf-8"
|
||||
except Exception:
|
||||
return "utf-8"
|
||||
return "utf-8"
|
||||
|
||||
|
||||
async def _standalone_read_text_limited(resp: Any, limit_bytes: int) -> str:
|
||||
body, _truncated = await _standalone_read_response_bytes_limited(resp, limit_bytes)
|
||||
if body is None:
|
||||
return await resp.text()
|
||||
return body.decode(_standalone_response_encoding(resp), "replace")
|
||||
|
||||
|
||||
async def _standalone_read_json_limited(resp: Any, limit_bytes: int) -> dict:
|
||||
body, truncated = await _standalone_read_response_bytes_limited(resp, limit_bytes)
|
||||
if body is None:
|
||||
return await resp.json()
|
||||
if truncated:
|
||||
raise ValueError(f"Discord API JSON response exceeds {limit_bytes} bytes")
|
||||
if not body:
|
||||
return {}
|
||||
data = json.loads(body.decode(_standalone_response_encoding(resp), "replace"))
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
async def _standalone_send(
|
||||
pconfig,
|
||||
chat_id: str,
|
||||
|
|
@ -7544,7 +7631,10 @@ async def _standalone_send(
|
|||
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=15), **_sess_kw) as info_sess:
|
||||
async with info_sess.get(info_url, headers=json_headers, **_req_kw) as info_resp:
|
||||
if info_resp.status == 200:
|
||||
info = await info_resp.json()
|
||||
info = await _standalone_read_json_limited(
|
||||
info_resp,
|
||||
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
|
||||
)
|
||||
is_forum = info.get("type") == 15
|
||||
_remember_channel_is_forum(chat_id, is_forum)
|
||||
except Exception:
|
||||
|
|
@ -7590,9 +7680,15 @@ async def _standalone_send(
|
|||
)
|
||||
async with session.post(thread_url, headers=auth_headers, data=form, **_req_kw) as resp:
|
||||
if resp.status not in {200, 201}:
|
||||
body = await resp.text()
|
||||
body = await _standalone_read_text_limited(
|
||||
resp,
|
||||
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
|
||||
)
|
||||
return {"error": f"Discord forum thread creation error ({resp.status}): {body}"}
|
||||
data = await resp.json()
|
||||
data = await _standalone_read_json_limited(
|
||||
resp,
|
||||
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
|
||||
)
|
||||
except Exception as e:
|
||||
return {"error": _standalone_sanitize_error(f"Discord forum thread upload failed: {e}")}
|
||||
else:
|
||||
|
|
@ -7608,9 +7704,15 @@ async def _standalone_send(
|
|||
**_req_kw,
|
||||
) as resp:
|
||||
if resp.status not in {200, 201}:
|
||||
body = await resp.text()
|
||||
body = await _standalone_read_text_limited(
|
||||
resp,
|
||||
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
|
||||
)
|
||||
return {"error": f"Discord forum thread creation error ({resp.status}): {body}"}
|
||||
data = await resp.json()
|
||||
data = await _standalone_read_json_limited(
|
||||
resp,
|
||||
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
|
||||
)
|
||||
|
||||
thread_id_created = data.get("id")
|
||||
starter_msg_id = (data.get("message") or {}).get("id", thread_id_created)
|
||||
|
|
@ -7632,9 +7734,15 @@ async def _standalone_send(
|
|||
if message.strip() or not media_files:
|
||||
async with session.post(url, headers=json_headers, json={"content": message}, **_req_kw) as resp:
|
||||
if resp.status not in {200, 201}:
|
||||
body = await resp.text()
|
||||
body = await _standalone_read_text_limited(
|
||||
resp,
|
||||
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
|
||||
)
|
||||
return {"error": f"Discord API error ({resp.status}): {body}"}
|
||||
last_data = await resp.json()
|
||||
last_data = await _standalone_read_json_limited(
|
||||
resp,
|
||||
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
|
||||
)
|
||||
|
||||
# Send each media file as a separate multipart upload
|
||||
for media_path, _is_voice in media_files:
|
||||
|
|
@ -7650,12 +7758,18 @@ async def _standalone_send(
|
|||
form.add_field("files[0]", f, filename=filename)
|
||||
async with session.post(url, headers=auth_headers, data=form, **_req_kw) as resp:
|
||||
if resp.status not in {200, 201}:
|
||||
body = await resp.text()
|
||||
body = await _standalone_read_text_limited(
|
||||
resp,
|
||||
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
|
||||
)
|
||||
warning = _standalone_sanitize_error(f"Failed to send media {media_path}: Discord API error ({resp.status}): {body}")
|
||||
logger.error(warning)
|
||||
warnings.append(warning)
|
||||
continue
|
||||
last_data = await resp.json()
|
||||
last_data = await _standalone_read_json_limited(
|
||||
resp,
|
||||
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
|
||||
)
|
||||
except Exception as e:
|
||||
warning = _standalone_sanitize_error(f"Failed to send media {media_path}: {e}")
|
||||
logger.error(warning)
|
||||
|
|
|
|||
|
|
@ -39,6 +39,8 @@ from tools.send_message_tool import (
|
|||
# and provide a thin ``_send_discord(token, ...)`` shim that mirrors the
|
||||
# pre-migration signature so the existing test bodies keep working.
|
||||
from plugins.platforms.discord.adapter import (
|
||||
_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES,
|
||||
_DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES,
|
||||
_derive_forum_thread_name,
|
||||
_probe_is_forum_cached,
|
||||
_remember_channel_is_forum,
|
||||
|
|
@ -68,6 +70,54 @@ async def _send_discord(
|
|||
)
|
||||
|
||||
|
||||
class _StreamingAiohttpContent:
|
||||
def __init__(self, body: bytes):
|
||||
self._body = body
|
||||
self.read_sizes = []
|
||||
|
||||
async def read(self, size=-1):
|
||||
self.read_sizes.append(size)
|
||||
if not self._body:
|
||||
return b""
|
||||
if size is None or size < 0:
|
||||
chunk = self._body
|
||||
self._body = b""
|
||||
return chunk
|
||||
chunk = self._body[:size]
|
||||
self._body = self._body[size:]
|
||||
return chunk
|
||||
|
||||
|
||||
class _StreamingAiohttpResponse:
|
||||
def __init__(self, status: int, body: bytes):
|
||||
self.status = status
|
||||
self.content = _StreamingAiohttpContent(body)
|
||||
self.closed = False
|
||||
self.json = AsyncMock(side_effect=AssertionError("resp.json() should not be used"))
|
||||
self.text = AsyncMock(side_effect=AssertionError("resp.text() should not be used"))
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _StreamingAiohttpSession:
|
||||
def __init__(self, response: _StreamingAiohttpResponse):
|
||||
self.response = response
|
||||
self.post = MagicMock(return_value=response)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
|
||||
|
||||
def _discord_entry():
|
||||
"""Return the live Discord PlatformEntry, importing lazily so plugin
|
||||
discovery is forced exactly once and patches survive across tests."""
|
||||
|
|
@ -1754,6 +1804,41 @@ class TestSendDiscordThreadId:
|
|||
assert "error" in result
|
||||
assert "403" in result["error"]
|
||||
|
||||
def test_success_response_json_read_is_bounded(self):
|
||||
"""Standalone Discord sends parse success JSON through the bounded reader."""
|
||||
body = b'{"id":"bounded-json"}'
|
||||
response = _StreamingAiohttpResponse(200, body)
|
||||
session = _StreamingAiohttpSession(response)
|
||||
|
||||
with patch("aiohttp.ClientSession", return_value=session):
|
||||
result = self._run("tok", "111", "hi", thread_id="999")
|
||||
|
||||
assert result["success"] is True
|
||||
assert result["message_id"] == "bounded-json"
|
||||
assert response.content.read_sizes[0] == _DISCORD_STANDALONE_JSON_BODY_LIMIT_BYTES + 1
|
||||
response.json.assert_not_awaited()
|
||||
response.text.assert_not_awaited()
|
||||
|
||||
def test_error_response_text_read_is_bounded(self):
|
||||
"""Oversized Discord API error bodies are capped before formatting."""
|
||||
body = b"E" * (_DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES + 1024)
|
||||
response = _StreamingAiohttpResponse(500, body)
|
||||
session = _StreamingAiohttpSession(response)
|
||||
|
||||
with patch("aiohttp.ClientSession", return_value=session):
|
||||
result = self._run("tok", "111", "hi", thread_id="999")
|
||||
|
||||
assert "error" in result
|
||||
assert "500" in result["error"]
|
||||
assert response.content.read_sizes[0] == _DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES + 1
|
||||
assert response.closed is True
|
||||
response.json.assert_not_awaited()
|
||||
response.text.assert_not_awaited()
|
||||
prefix = "Discord API error (500): "
|
||||
assert len(result["error"].encode("utf-8")) <= (
|
||||
len(prefix.encode("utf-8")) + _DISCORD_STANDALONE_ERROR_BODY_LIMIT_BYTES
|
||||
)
|
||||
|
||||
|
||||
class TestSendToPlatformDiscordThread:
|
||||
"""_send_to_platform passes thread_id through to _send_discord."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue