fix(channels): offload outbound attachment file IO (#4633)

This commit is contained in:
Felix Wang 2026-08-02 09:08:36 +08:00 committed by GitHub
parent cd8825b0f0
commit 4795452102
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 239 additions and 26 deletions

View File

@ -267,6 +267,8 @@ Blocking-IO runtime gate (`tests/blocking_io/`):
`test_feishu_receive_file.py` (locks Feishu attachment path preparation and `test_feishu_receive_file.py` (locks Feishu attachment path preparation and
persistence plus remote sandbox acquisition/sync off the event loop, and persistence plus remote sandbox acquisition/sync off the event loop, and
skips redundant sandbox sync when thread data is already mounted); skips redundant sandbox sync when thread data is already mounted);
`test_channel_outbound_files.py` (locks Feishu, Telegram, and WeCom outbound
attachment open/read/hash work off the event loop);
`test_openviking_memory_backend.py` (locks the OpenViking backend's async `test_openviking_memory_backend.py` (locks the OpenViking backend's async
add/context/search entrypoints offloading synchronous HTTP and watermark add/context/search entrypoints offloading synchronous HTTP and watermark
filesystem IO); and filesystem IO); and

View File

@ -325,15 +325,23 @@ class FeishuChannel(Channel):
logger.exception("[Feishu] failed to upload/send file: %s", attachment.filename) logger.exception("[Feishu] failed to upload/send file: %s", attachment.filename)
return False return False
async def _upload_image(self, path) -> str: def _upload_image_sync(self, path):
"""Upload an image to Feishu and return the image_key."""
with open(str(path), "rb") as f: with open(str(path), "rb") as f:
request = self._CreateImageRequest.builder().request_body(self._CreateImageRequestBody.builder().image_type("message").image(f).build()).build() request = self._CreateImageRequest.builder().request_body(self._CreateImageRequestBody.builder().image_type("message").image(f).build()).build()
response = await asyncio.to_thread(self._api_client.im.v1.image.create, request) return self._api_client.im.v1.image.create(request)
async def _upload_image(self, path) -> str:
"""Upload an image to Feishu and return the image_key."""
response = await asyncio.to_thread(self._upload_image_sync, path)
if not response.success(): if not response.success():
raise RuntimeError(f"Feishu image upload failed: code={response.code}, msg={response.msg}") raise RuntimeError(f"Feishu image upload failed: code={response.code}, msg={response.msg}")
return response.data.image_key return response.data.image_key
def _upload_file_sync(self, path, filename: str, file_type: str):
with open(str(path), "rb") as f:
request = self._CreateFileRequest.builder().request_body(self._CreateFileRequestBody.builder().file_type(file_type).file_name(filename).file(f).build()).build()
return self._api_client.im.v1.file.create(request)
async def _upload_file(self, path, filename: str) -> str: async def _upload_file(self, path, filename: str) -> str:
"""Upload a file to Feishu and return the file_key.""" """Upload a file to Feishu and return the file_key."""
suffix = path.suffix.lower() if hasattr(path, "suffix") else "" suffix = path.suffix.lower() if hasattr(path, "suffix") else ""
@ -348,9 +356,7 @@ class FeishuChannel(Channel):
else: else:
file_type = "stream" file_type = "stream"
with open(str(path), "rb") as f: response = await asyncio.to_thread(self._upload_file_sync, path, filename, file_type)
request = self._CreateFileRequest.builder().request_body(self._CreateFileRequestBody.builder().file_type(file_type).file_name(filename).file(f).build()).build()
response = await asyncio.to_thread(self._api_client.im.v1.file.create, request)
if not response.success(): if not response.success():
raise RuntimeError(f"Feishu file upload failed: code={response.code}, msg={response.msg}") raise RuntimeError(f"Feishu file upload failed: code={response.code}, msg={response.msg}")
return response.data.file_key return response.data.file_key

View File

@ -42,6 +42,12 @@ MAX_TRACKED_STREAM_MESSAGES = 256
_monotonic = time.monotonic _monotonic = time.monotonic
def _load_telegram_input_file(path, filename: str):
from telegram import InputFile
return InputFile(path.read_bytes(), filename=filename)
class TelegramChannel(Channel): class TelegramChannel(Channel):
"""Telegram bot channel using long-polling. """Telegram bot channel using long-polling.
@ -385,21 +391,17 @@ class TelegramChannel(Channel):
reply_to = self._last_bot_message.get(msg.chat_id) reply_to = self._last_bot_message.get(msg.chat_id)
try: try:
input_file = await asyncio.to_thread(_load_telegram_input_file, attachment.actual_path, attachment.filename)
if attachment.is_image and attachment.size <= 10 * 1024 * 1024: if attachment.is_image and attachment.size <= 10 * 1024 * 1024:
with open(attachment.actual_path, "rb") as f: kwargs: dict[str, Any] = {"chat_id": chat_id, "photo": input_file}
kwargs: dict[str, Any] = {"chat_id": chat_id, "photo": f} if reply_to:
if reply_to: kwargs["reply_to_message_id"] = reply_to
kwargs["reply_to_message_id"] = reply_to sent = await bot.send_photo(**kwargs)
sent = await bot.send_photo(**kwargs)
else: else:
from telegram import InputFile kwargs = {"chat_id": chat_id, "document": input_file}
if reply_to:
with open(attachment.actual_path, "rb") as f: kwargs["reply_to_message_id"] = reply_to
input_file = InputFile(f, filename=attachment.filename) sent = await bot.send_document(**kwargs)
kwargs = {"chat_id": chat_id, "document": input_file}
if reply_to:
kwargs["reply_to_message_id"] = reply_to
sent = await bot.send_document(**kwargs)
self._last_bot_message[msg.chat_id] = sent.message_id self._last_bot_message[msg.chat_id] = sent.message_id
logger.info("[Telegram] file sent: %s to chat=%s", attachment.filename, msg.chat_id) logger.info("[Telegram] file sent: %s to chat=%s", attachment.filename, msg.chat_id)

View File

@ -21,6 +21,18 @@ from app.channels.message_bus import (
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def _file_md5(path: str) -> str:
md5_hasher = hashlib.md5()
with open(path, "rb") as file_obj:
for chunk in iter(lambda: file_obj.read(1024 * 1024), b""):
md5_hasher.update(chunk)
return md5_hasher.hexdigest()
def _open_binary(path: str):
return open(path, "rb")
class WeComChannel(Channel): class WeComChannel(Channel):
def __init__(self, bus: MessageBus, config: dict[str, Any]) -> None: def __init__(self, bus: MessageBus, config: dict[str, Any]) -> None:
super().__init__(name="wecom", bus=bus, config=config) super().__init__(name="wecom", bus=bus, config=config)
@ -428,11 +440,7 @@ class WeComChannel(Channel):
logger.warning("[WeCom] invalid total_chunks=%d for %s", total_chunks, filename) logger.warning("[WeCom] invalid total_chunks=%d for %s", total_chunks, filename)
return None return None
md5_hasher = hashlib.md5() md5 = await asyncio.to_thread(_file_md5, path)
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(1024 * 1024), b""):
md5_hasher.update(chunk)
md5 = md5_hasher.hexdigest()
init_req_id = generate_req_id("aibot_upload_media_init") init_req_id = generate_req_id("aibot_upload_media_init")
init_body = { init_body = {
@ -448,9 +456,10 @@ class WeComChannel(Channel):
logger.warning("[WeCom] upload init returned no upload_id: %s", init_ack) logger.warning("[WeCom] upload init returned no upload_id: %s", init_ack)
return None return None
with open(path, "rb") as f: file_obj = await asyncio.to_thread(_open_binary, path)
try:
for idx in range(total_chunks): for idx in range(total_chunks):
data = f.read(chunk_size) data = await asyncio.to_thread(file_obj.read, chunk_size)
if not data: if not data:
break break
chunk_req_id = generate_req_id("aibot_upload_media_chunk") chunk_req_id = generate_req_id("aibot_upload_media_chunk")
@ -460,6 +469,8 @@ class WeComChannel(Channel):
"base64_data": base64.b64encode(data).decode("utf-8"), "base64_data": base64.b64encode(data).decode("utf-8"),
} }
await self._send_ws_upload_command(chunk_req_id, chunk_body, "aibot_upload_media_chunk") await self._send_ws_upload_command(chunk_req_id, chunk_body, "aibot_upload_media_chunk")
finally:
await asyncio.to_thread(file_obj.close)
finish_req_id = generate_req_id("aibot_upload_media_finish") finish_req_id = generate_req_id("aibot_upload_media_finish")
finish_ack = await self._send_ws_upload_command(finish_req_id, {"upload_id": upload_id}, "aibot_upload_media_finish") finish_ack = await self._send_ws_upload_command(finish_req_id, {"upload_id": upload_id}, "aibot_upload_media_finish")

View File

@ -0,0 +1,192 @@
"""Regression anchors for outbound IM attachment file IO.
Feishu, Telegram, and WeCom send attachments from async channel handlers. File
open/read/hash work must run off the event loop; otherwise a large outbound
artifact stalls every channel and Gateway coroutine on that worker.
"""
from __future__ import annotations
import asyncio
import base64
import builtins
import hashlib
import threading
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from app.channels.feishu import FeishuChannel
from app.channels.message_bus import MessageBus, OutboundMessage, ResolvedAttachment
from app.channels.telegram import TelegramChannel
from app.channels.wecom import WeComChannel
pytestmark = pytest.mark.asyncio
def _attachment(path: Path, *, is_image: bool) -> ResolvedAttachment:
return ResolvedAttachment(
virtual_path=f"/mnt/user-data/outputs/{path.name}",
actual_path=path,
filename=path.name,
mime_type="image/png" if is_image else "application/octet-stream",
size=path.stat().st_size,
is_image=is_image,
)
def _outbound(channel_name: str) -> OutboundMessage:
return OutboundMessage(
channel_name=channel_name,
chat_id="123",
thread_id="thread-1",
text="attachment",
)
def _builder(*, captured_file: dict[str, object] | None = None) -> MagicMock:
builder = MagicMock()
for method_name in ("request_body", "image_type", "file_type", "file_name"):
getattr(builder, method_name).return_value = builder
if captured_file is not None:
def _capture(file_obj):
captured_file["file"] = file_obj
return builder
builder.image.side_effect = _capture
builder.file.side_effect = _capture
builder.build.return_value = object()
return builder
async def test_feishu_outbound_uploads_do_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
image_path = tmp_path / "chart.png"
file_path = tmp_path / "report.pdf"
await asyncio.to_thread(image_path.write_bytes, b"image-bytes")
await asyncio.to_thread(file_path.write_bytes, b"file-bytes")
channel = FeishuChannel(MessageBus(), {})
channel._api_client = MagicMock()
real_open = builtins.open
opened_on_threads: list[int] = []
tracked_paths = {str(image_path), str(file_path)}
def _tracked_open(file, *args, **kwargs):
if str(file) in tracked_paths:
opened_on_threads.append(threading.get_ident())
return real_open(file, *args, **kwargs)
monkeypatch.setattr(builtins, "open", _tracked_open)
image_capture: dict[str, object] = {}
channel._CreateImageRequest = MagicMock()
channel._CreateImageRequest.builder.return_value = _builder()
channel._CreateImageRequestBody = MagicMock()
channel._CreateImageRequestBody.builder.return_value = _builder(captured_file=image_capture)
image_response = MagicMock()
image_response.success.return_value = True
image_response.data.image_key = "image-key"
image_worker_thread: int | None = None
def _create_image(_request):
nonlocal image_worker_thread
image_worker_thread = threading.get_ident()
file_obj = image_capture["file"]
assert not file_obj.closed
assert file_obj.read() == b"image-bytes"
return image_response
channel._api_client.im.v1.image.create.side_effect = _create_image
file_capture: dict[str, object] = {}
channel._CreateFileRequest = MagicMock()
channel._CreateFileRequest.builder.return_value = _builder()
channel._CreateFileRequestBody = MagicMock()
channel._CreateFileRequestBody.builder.return_value = _builder(captured_file=file_capture)
file_response = MagicMock()
file_response.success.return_value = True
file_response.data.file_key = "file-key"
file_worker_thread: int | None = None
def _create_file(_request):
nonlocal file_worker_thread
file_worker_thread = threading.get_ident()
file_obj = file_capture["file"]
assert not file_obj.closed
assert file_obj.read() == b"file-bytes"
return file_response
channel._api_client.im.v1.file.create.side_effect = _create_file
event_loop_thread = threading.get_ident()
assert await channel._upload_image(image_path) == "image-key"
assert await channel._upload_file(file_path, file_path.name) == "file-key"
assert len(opened_on_threads) == 2
assert all(thread_id != event_loop_thread for thread_id in opened_on_threads)
assert image_worker_thread != event_loop_thread
assert file_worker_thread != event_loop_thread
@pytest.mark.parametrize("is_image", [False, True])
async def test_telegram_outbound_upload_does_not_block_event_loop(tmp_path: Path, is_image: bool) -> None:
path = tmp_path / ("chart.png" if is_image else "report.bin")
payload = b"telegram-payload"
await asyncio.to_thread(path.write_bytes, payload)
sent_file = None
class _Bot:
async def send_document(self, **kwargs):
nonlocal sent_file
sent_file = kwargs["document"]
return SimpleNamespace(message_id=7)
async def send_photo(self, **kwargs):
nonlocal sent_file
sent_file = kwargs["photo"]
if hasattr(sent_file, "read"):
sent_file.read()
return SimpleNamespace(message_id=7)
channel = TelegramChannel(MessageBus(), {})
channel._application = SimpleNamespace(bot=_Bot())
assert await channel.send_file(_outbound("telegram"), _attachment(path, is_image=is_image))
assert sent_file.input_file_content == payload
assert sent_file.filename == path.name
async def test_wecom_outbound_upload_does_not_block_event_loop(tmp_path: Path) -> None:
path = tmp_path / "report.bin"
payload = b"x" * (512 * 1024 + 17)
await asyncio.to_thread(path.write_bytes, payload)
channel = WeComChannel(MessageBus(), {})
channel._ws_client = object()
channel._send_ws_upload_command = AsyncMock(
side_effect=[
{"body": {"upload_id": "upload-1"}},
{"body": {}},
{"body": {}},
{"body": {"media_id": "media-1"}},
]
)
result = await channel._upload_media_ws(
media_type="file",
filename=path.name,
path=str(path),
size=len(payload),
)
assert result == "media-1"
calls = channel._send_ws_upload_command.await_args_list
assert calls[0].args[1]["md5"] == hashlib.md5(payload).hexdigest()
encoded_chunks = [call.args[1]["base64_data"] for call in calls[1:-1]]
assert b"".join(base64.b64decode(chunk) for chunk in encoded_chunks) == payload