ouyangbo/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

refactor(gateway/yuanbao): fold GroupQueryService into adapter, drop sender hook indirection, _media_ref/_cancel_task helpers

Browse Source
This commit is contained in:
Teknium
2026-09-02 21:36:31 -07:00
parent 445f932325
commit 43159ca0c8
1 changed files with 105 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gateway/platforms/yuanbao.py
+105 -151
View File
@@ -130,6 +130,15 @@ def _cancel_all(tasks: Dict[str, asyncio.Task]) -> None:
tasks.clear()
async def _cancel_task(task: asyncio.Task) -> None:
"""Cancel *task* and wait for it to unwind (swallowing the CancelledError)."""
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
class MarkdownProcessor:
"""Thin delegates to the shared fence-aware chunker in gateway.platforms.helpers; method names
kept for existing call sites and tests."""
@@ -306,9 +315,7 @@ class InboundContext:
reply_to_message_id: Optional[str] = None
reply_to_text: Optional[str] = None
quote_media_refs: list = dc_field(default_factory=list) # (rid, kind, filename)
# MediaResolveMiddleware: resolved local paths, deduped, in this order: 1) media carried by
# the current message, 2) quoted-message media (when reply_to_message_id is set), 3) recent
# group-observed media (group chats without a quote only).
# MediaResolveMiddleware: deduped local paths — own media, then quoted media, else group-observed media
media_urls: list = dc_field(default_factory=list)
media_types: list = dc_field(default_factory=list)
channel_prompt: Optional[str] = None # GroupAttributionMiddleware
@@ -470,32 +477,25 @@ class DecodeMiddleware(InboundMiddleware):
async def handle(self, ctx: InboundContext, next_fn) -> None:
if not ctx.raw_frames:
return # Stop pipeline — nothing to decode
merged_push = None
decoded_via = ""
merged: Optional[dict] = None
for data in ctx.raw_frames:
push, via = self._decode_single(ctx.adapter, data)
if not push:
logger.info("[%s] Push decoded but no valid message. raw hex(first64)=%s",
ctx.adapter.name, data.hex()[:128] if data else "(empty)")
continue
if merged_push is None:
merged_push = push
decoded_via = via
elif merged is None:
merged, ctx.decoded_via = push, via
logger.info("[%s] Frame decoded (via=%s): len=%d", ctx.adapter.name, via, len(data))
else:
# Subsequent pushes: append msg_body to the base, newline-separated.
extra_body = push.get("msg_body", [])
if extra_body:
merged_push["msg_body"] = merged_push.get("msg_body", []) + [_text_elem("\n")] + extra_body
logger.info("[%s] Merged %d extra msg_body elements from aggregated push", ctx.adapter.name, len(extra_body))
if not merged_push:
elif push.get("msg_body", []): # subsequent pushes: append msg_body, newline-separated
merged["msg_body"] = merged.get("msg_body", []) + [_text_elem("\n")] + push["msg_body"]
logger.info("[%s] Merged %d extra msg_body elements from aggregated push", ctx.adapter.name, len(push["msg_body"]))
if not merged:
return # Stop pipeline
ctx.push = merged_push
ctx.decoded_via = decoded_via
ctx.push = merged
logger.info(
"[%s] Push decoded (via=%s): from=%s group=%s msg_id=%s msg_types=%s",
ctx.adapter.name, ctx.decoded_via, ctx.push.get("from_account", ""), ctx.push.get("group_code", ""),
ctx.push.get("msg_id", ""), [e.get("msg_type", "") for e in ctx.push.get("msg_body", [])],
ctx.adapter.name, ctx.decoded_via, merged.get("from_account", ""), merged.get("group_code", ""),
merged.get("msg_id", ""), [e.get("msg_type", "") for e in merged.get("msg_body", [])],
)
logger.debug("[%s] Push payload: %s", ctx.adapter.name, ctx.push)
await next_fn()
@@ -562,8 +562,7 @@ class RecallGuardMiddleware(InboundMiddleware):
if cmd == "Group.CallbackAfterRecallMsg":
seq_list = push.get("recall_msg_seq_list") or []
else:
mid = push.get("msg_id") or ""
seq = push.get("msg_seq")
mid, seq = push.get("msg_id") or "", push.get("msg_seq")
seq_list = [{"msg_id": mid, "msg_seq": seq}] if (mid or seq) else []
if not seq_list:
logger.debug("[%s] Recall callback with empty seq_list, skipping", adapter.name)
@@ -823,6 +822,14 @@ def _file_name(content: dict) -> str:
or str(content.get("filename") or "").strip())
def _media_ref(kind: str, url: str, name: str = "") -> Dict[str, str]:
"""media_refs entry; ``name`` is only emitted when non-empty."""
ref: Dict[str, str] = {"kind": kind, "url": url}
if name:
ref["name"] = name
return ref
class ExtractContentMiddleware(InboundMiddleware):
"""Extract raw text, media refs and forwarded records from msg_body."""
name = "extract-content"
@@ -969,14 +976,11 @@ class ExtractContentMiddleware(InboundMiddleware):
if msg_type == "TIMImageElem":
image_url = ExtractContentMiddleware._pick_image_url(content)
if image_url:
refs.append({"kind": "image", "url": image_url})
refs.append(_media_ref("image", image_url))
elif msg_type == "TIMFileElem":
file_url = str(content.get("url") or "").strip()
if file_url:
ref: Dict[str, str] = {"kind": "file", "url": file_url}
if _file_name(content):
ref["name"] = _file_name(content)
refs.append(ref)
refs.append(_media_ref("file", file_url, _file_name(content)))
return refs
@staticmethod
@@ -1296,23 +1300,17 @@ class ForwardedRecordsParseMiddleware(InboundMiddleware):
file_name = str(media.get("file_name") or "").strip()
# media_id is directly usable as a ybres RID; else parse resourceId from the URL.
rid = str(media.get("media_id") or "").strip() or ExtractContentMiddleware._parse_resource_id(url)
if media_type == "image":
if url and rid:
return f"[image|ybres:{rid}] {file_name}".rstrip(), {"kind": "image", "url": url}
kind = {"image": "image", "file": "file", "document": "file", "code": "file", "video": "video"}.get(media_type)
if kind and url and rid:
return f"[{kind}|ybres:{rid}] {file_name}".rstrip(), _media_ref(kind, url, file_name if kind == "file" else "")
if kind == "image":
return f"[image] {file_name or plain_text}".rstrip(), None
if media_type in ("file", "document", "code"):
if url and rid:
ref: Dict[str, str] = {"kind": "file", "url": url}
if file_name:
ref["name"] = file_name
return f"[file|ybres:{rid}] {file_name}".rstrip(), ref
if kind == "file":
return f"[file] {file_name}".rstrip(), None
if kind == "video":
return f"[video] {file_name or url}".rstrip(), None
if media_type == "url": # link share (e.g. WeChat article) — keep URL for the agent
return f"[link] {file_name or str(media.get('title') or '')} {url}".rstrip(), None
if media_type == "video":
if url and rid:
return f"[video|ybres:{rid}] {file_name}".rstrip(), {"kind": "video", "url": url}
return f"[video] {file_name or url}".rstrip(), None
return f"[{media_type or 'media'}] {url or file_name}".rstrip(), None
@classmethod
@@ -1946,11 +1944,7 @@ class ConnectionManager:
for attr, _coro_name, _tag in self._LOOPS:
task = getattr(self, attr)
if task:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
await _cancel_task(task)
setattr(self, attr, None)
disc_exc = RuntimeError("YuanbaoAdapter disconnected")
for fut in self._pending_acks.values():
@@ -1974,10 +1968,9 @@ class ConnectionManager:
))
logger.debug("[%s] AUTH_BIND sent (msg_id=%s uid=%s)", adapter.name, msg_id, uid)
try:
_loop = asyncio.get_running_loop()
deadline = _loop.time() + AUTH_TIMEOUT_SECONDS
deadline = asyncio.get_running_loop().time() + AUTH_TIMEOUT_SECONDS
while True:
remaining = deadline - _loop.time()
remaining = deadline - asyncio.get_running_loop().time()
if remaining <= 0:
logger.error("[%s] AUTH_BIND timeout waiting for BIND_ACK", adapter.name)
return False
@@ -1989,14 +1982,14 @@ class ConnectionManager:
except Exception:
continue
head = msg.get("head", {})
if head.get("cmd_type", -1) == CMD_TYPE["Response"] and head.get("cmd", "") == "auth-bind":
connect_id = self._extract_connect_id(msg)
if not connect_id:
logger.error("[%s] BIND_ACK missing connectId", adapter.name)
return False
self._connect_id = connect_id
logger.info("[%s] BIND_ACK received: connectId=%s", adapter.name, connect_id)
return True
if head.get("cmd_type", -1) != CMD_TYPE["Response"] or head.get("cmd", "") != "auth-bind":
continue
self._connect_id = self._extract_connect_id(msg)
if not self._connect_id:
logger.error("[%s] BIND_ACK missing connectId", adapter.name)
return False
logger.info("[%s] BIND_ACK received: connectId=%s", adapter.name, self._connect_id)
return True
except asyncio.TimeoutError:
logger.error("[%s] AUTH_BIND timeout", adapter.name)
except Exception as exc:
@@ -2094,12 +2087,9 @@ class ConnectionManager:
if cmd_type == CMD_TYPE["Response"]:
if cmd == "ping": # HEARTBEAT_ACK
logger.debug("[%s] HEARTBEAT_ACK received (msg_id=%s)", adapter.name, msg_id)
if self._pending_pong is not None and not self._pending_pong.done():
self._pending_pong.set_result(True)
else:
fut = self._pop_pending(msg_id)
if fut is not None:
fut.set_result(True)
pong = self._pending_pong if self._pending_pong is not None and not self._pending_pong.done() else self._pop_pending(msg_id)
if pong is not None:
pong.set_result(True)
elif cmd in {"send_group_heartbeat", "send_private_heartbeat"}:
# Fire-and-forget heartbeat ACKs: nobody awaits them; discard to avoid "Unmatched" noise.
logger.debug("[%s] Heartbeat ACK received: cmd=%s msg_id=%s", adapter.name, cmd, msg_id)
@@ -2286,18 +2276,15 @@ class MediaSendHandler(ABC):
token_data = await adapter._get_cached_token()
credentials = await get_cos_credentials(
app_key=adapter._app_key, api_domain=adapter._api_domain, token=token_data.get("token", ""),
filename=filename, bot_id=token_data.get("bot_id", "") or adapter._bot_id or "",
route_env=adapter._route_env,
filename=filename, bot_id=token_data.get("bot_id", "") or adapter._bot_id or "", route_env=adapter._route_env,
)
upload_result = await upload_to_cos(
file_bytes=file_bytes, filename=filename, content_type=content_type, credentials=credentials,
bucket=credentials["bucketName"], region=credentials["region"],
)
# Drop keys passed explicitly below to avoid "multiple values" TypeError.
fwd_kwargs = {k: v for k, v in kwargs.items() if k not in {"file_uuid", "filename", "content_type"}}
msg_body = self.build_msg_body(
upload_result, file_uuid=md5_hex(file_bytes), filename=filename, content_type=content_type, **fwd_kwargs,
)
# Explicit keys win over caller kwargs (avoids "multiple values" TypeError).
msg_body = self.build_msg_body(upload_result, **{
**kwargs, "file_uuid": md5_hex(file_bytes), "filename": filename, "content_type": content_type})
else:
msg_body = self.build_msg_body({}, **kwargs)
if caption:
@@ -2370,49 +2357,6 @@ class StickerHandler(MediaSendHandler):
return build_sticker_msg_body(get_random_sticker())
class GroupQueryService:
"""Low-level WS group queries (group info, member list); populates adapter._member_cache."""
def __init__(self, adapter: "YuanbaoAdapter") -> None:
self._adapter = adapter
async def _ws_query(self, label: str, group_code: str, encoded: bytes, decode_rsp, empty: dict) -> Optional[dict]:
"""Send an encoded group query over WS; return decoded biz payload, *empty* when none, None on failure."""
adapter = self._adapter
if adapter._connection.ws is None:
return None
req_id = decode_conn_msg(encoded)["head"]["msg_id"]
try:
response = await adapter._connection.send_biz_request(encoded, req_id=req_id)
status = response.get("head", {}).get("status", 0)
if status != 0:
logger.warning("[%s] %s failed: status=%d", adapter.name, label, status)
return None
biz_data = response.get("data", b"") or response.get("body", b"")
return decode_rsp(biz_data) if biz_data and isinstance(biz_data, bytes) else empty
except asyncio.TimeoutError:
logger.warning("[%s] %s timeout: group=%s", adapter.name, label, group_code)
except Exception as exc:
logger.warning("[%s] %s failed: %s", adapter.name, label, exc)
return None
async def query_group_info_raw(self, group_code: str) -> Optional[dict]:
"""Group info (name, owner, member count…); None on failure."""
return await self._ws_query(
"query_group_info", group_code, encode_query_group_info(group_code),
decode_query_group_info_rsp, {"group_code": group_code},
)
async def get_group_member_list_raw(self, group_code: str, offset: int = 0, limit: int = 200) -> Optional[dict]:
"""Group member list; None on failure. Populates adapter._member_cache."""
result = await self._ws_query(
"get_group_member_list", group_code, encode_get_group_member_list(group_code, offset=offset, limit=limit),
decode_get_group_member_list_rsp, {"members": [], "next_offset": 0, "is_complete": True},
)
if result and result.get("members"):
self._adapter._member_cache[group_code] = (time.time(), result["members"])
return result
class HeartbeatManager:
"""Reply heartbeat lifecycle: RUNNING every 2s, auto-FINISH after 30s idle, explicit stop."""
def __init__(self, adapter: "YuanbaoAdapter") -> None:
@@ -2420,18 +2364,20 @@ class HeartbeatManager:
self._reply_heartbeat_tasks: Dict[str, asyncio.Task] = {}
self._reply_hb_last_active: Dict[str, float] = {}
def _ready(self) -> bool:
return self._adapter._connection.ws is not None and bool(self._adapter._bot_id)
async def send_heartbeat_once(self, chat_id: str, heartbeat_val: int) -> None:
"""Send a single heartbeat (RUNNING or FINISH), best effort."""
adapter = self._adapter
conn = adapter._connection
if conn.ws is None or not adapter._bot_id:
if not self._ready():
return
try:
if chat_id.startswith("group:"):
encoded = encode_send_group_heartbeat(from_account=adapter._bot_id, group_code=chat_id[len("group:"):], heartbeat=heartbeat_val)
else:
encoded = encode_send_private_heartbeat(from_account=adapter._bot_id, to_account=chat_id.removeprefix("direct:"), heartbeat=heartbeat_val)
await conn.ws.send(encoded)
await adapter._connection.ws.send(encoded)
logger.debug("[%s] Reply heartbeat %s sent: chat=%s", adapter.name,
"RUNNING" if heartbeat_val == WS_HEARTBEAT_RUNNING else "FINISH", chat_id)
except Exception as exc:
@@ -2439,8 +2385,7 @@ class HeartbeatManager:
async def start(self, chat_id: str) -> None:
"""Start or renew the periodic RUNNING sender."""
adapter = self._adapter
if adapter._connection.ws is None or not adapter._bot_id:
if not self._ready():
return
self._reply_hb_last_active[chat_id] = time.time()
existing = self._reply_heartbeat_tasks.get(chat_id)
@@ -2475,11 +2420,7 @@ class HeartbeatManager:
"""Stop the RUNNING sender and optionally send FINISH."""
task = self._reply_heartbeat_tasks.pop(chat_id, None)
if task and not task.done():
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
await _cancel_task(task)
if send_finish:
await self.send_heartbeat_once(chat_id, WS_HEARTBEAT_FINISH)
@@ -2529,9 +2470,6 @@ class MessageSender:
def __init__(self, adapter: "YuanbaoAdapter") -> None:
self._adapter = adapter
self._chat_locks: collections.OrderedDict[str, asyncio.Lock] = collections.OrderedDict()
# Hooks injected by OutboundManager: cancel slow-notifier / send FINISH heartbeat
self._on_send_start: Optional[Callable[[str], Any]] = None
self._on_send_finish: Optional[Callable[[str], Any]] = None
self._media_handlers: Dict[str, MediaSendHandler] = {
"image_url": ImageUrlHandler(), "image_file": ImageFileHandler(), "document": DocumentHandler(), "sticker": StickerHandler(),
}
@@ -2552,8 +2490,7 @@ class MessageSender:
adapter = self._adapter
if adapter._connection.ws is None:
return SendResult(success=False, error="Not connected", retryable=True)
if self._on_send_start:
self._on_send_start(chat_id)
adapter._outbound.cancel_slow_notifier(chat_id)
async with self.get_chat_lock(chat_id):
content_to_send = self.strip_cron_wrapper(content)
chunks = self.truncate_message(content_to_send, adapter.MAX_TEXT_CHUNK)
@@ -2563,11 +2500,10 @@ class MessageSender:
result = await self.send_text_chunk(chat_id, chunk, reply_to if i == 0 else None, group_code=group_code)
if not result.success:
return result
if self._on_send_finish: # e.g. FINISH heartbeat
try:
await self._on_send_finish(chat_id)
except Exception:
pass
try: # delivery done → FINISH heartbeat (ordering: RUNNING… → message → FINISH)
await adapter._outbound.heartbeat.send_heartbeat_once(chat_id, WS_HEARTBEAT_FINISH)
except Exception:
pass
return SendResult(success=True)
async def send_media(self, chat_id: str, handler_name: str, reply_to: Optional[str] = None,
@@ -2732,7 +2668,8 @@ class MessageSender:
class OutboundManager:
"""Composes MessageSender, HeartbeatManager and SlowResponseNotifier and wires their coordination hooks."""
"""Composes MessageSender, HeartbeatManager and SlowResponseNotifier (sender cancels the notifier
before a send and emits the FINISH heartbeat after)."""
CHAT_DICT_MAX_SIZE: ClassVar[int] = MessageSender.CHAT_DICT_MAX_SIZE
def __init__(self, adapter: "YuanbaoAdapter") -> None:
@@ -2740,12 +2677,6 @@ class OutboundManager:
self.sender: MessageSender = MessageSender(adapter)
self.heartbeat: HeartbeatManager = HeartbeatManager(adapter)
self.slow_notifier: SlowResponseNotifier = SlowResponseNotifier(adapter, self.sender)
self.sender._on_send_start = self.slow_notifier.cancel # before a send: drop the slow-response notice
self.sender._on_send_finish = self._handle_send_finish
async def _handle_send_finish(self, chat_id: str) -> None:
"""After a send: FINISH heartbeat."""
await self.heartbeat.send_heartbeat_once(chat_id, WS_HEARTBEAT_FINISH)
async def start_slow_notifier(self, chat_id: str) -> None:
await self.slow_notifier.start(chat_id)
@@ -2822,7 +2753,6 @@ class YuanbaoAdapter(BasePlatformAdapter):
group_policy, group_allow_from = _policy("group")
self._access_policy = AccessPolicy(dm_policy=dm_policy, dm_allow_from=dm_allow_from,
group_policy=group_policy, group_allow_from=group_allow_from)
self._group_query = GroupQueryService(self)
self._inbound_pipeline: InboundPipeline = InboundPipelineBuilder.build()
# Auto-sethome stays open when no home is set or the home is a group (upgradable by first DM).
_existing_home = os.getenv("YUANBAO_HOME_CHANNEL") or (config.home_channel.chat_id if config.home_channel else "")
@@ -2853,15 +2783,12 @@ class YuanbaoAdapter(BasePlatformAdapter):
from gateway.pairing import PairingStore
return PairingStore().is_approved(Platform.YUANBAO.value, sender)
return False
if ctx.chat_type == "group":
group_code = str(ctx.group_code or "").strip()
if not group_code:
return False
if policy.group_policy == "allowlist":
return policy.is_group_allowed(group_code)
if policy.group_policy == "open":
return policy._open_dm_opted_in()
return False
group_code = str(ctx.group_code or "").strip()
if ctx.chat_type != "group" or not group_code:
return False
if policy.group_policy == "allowlist":
return policy.is_group_allowed(group_code)
return policy.group_policy == "open" and policy._open_dm_opted_in()
async def connect(self, *, is_reconnect: bool = False) -> bool:
ok = await self._connection.open()
@@ -2923,11 +2850,38 @@ class YuanbaoAdapter(BasePlatformAdapter):
self._processing_msg_ids.pop(session_key, None)
self._processing_msg_texts.pop(session_key, None)
async def _ws_query(self, label: str, group_code: str, encoded: bytes, decode_rsp, empty: dict) -> Optional[dict]:
"""Send an encoded group query over WS; return decoded biz payload, *empty* when none, None on failure."""
if self._connection.ws is None:
return None
try:
response = await self._connection.send_biz_request(encoded, req_id=decode_conn_msg(encoded)["head"]["msg_id"])
status = response.get("head", {}).get("status", 0)
if status != 0:
logger.warning("[%s] %s failed: status=%d", self.name, label, status)
return None
biz_data = response.get("data", b"") or response.get("body", b"")
return decode_rsp(biz_data) if biz_data and isinstance(biz_data, bytes) else empty
except asyncio.TimeoutError:
logger.warning("[%s] %s timeout: group=%s", self.name, label, group_code)
except Exception as exc:
logger.warning("[%s] %s failed: %s", self.name, label, exc)
return None
async def query_group_info(self, group_code: str) -> Optional[dict]:
return await self._group_query.query_group_info_raw(group_code)
"""Group info (name, owner, member count…); None on failure."""
return await self._ws_query("query_group_info", group_code, encode_query_group_info(group_code),
decode_query_group_info_rsp, {"group_code": group_code})
async def get_group_member_list(self, group_code: str, offset: int = 0, limit: int = 200) -> Optional[dict]:
return await self._group_query.get_group_member_list_raw(group_code, offset=offset, limit=limit)
"""Group member list; None on failure. Populates ``_member_cache`` for @mention resolution."""
result = await self._ws_query(
"get_group_member_list", group_code, encode_get_group_member_list(group_code, offset=offset, limit=limit),
decode_get_group_member_list_rsp, {"members": [], "next_offset": 0, "is_complete": True},
)
if result and result.get("members"):
self._member_cache[group_code] = (time.time(), result["members"])
return result
async def send_dm(self, user_id: str, text: str, group_code: str = "") -> SendResult:
"""Proactive C2C DM (text capped at DM_MAX_CHARS); group_code marks a group-originated DM."""