""" Yuanbao platform adapter. Connects to the Yuanbao WebSocket gateway, handles authentication (AUTH_BIND), heartbeat, reconnection, message receive (T05) and send (T06). Configuration in config.yaml (or via env vars): platforms: yuanbao: extra: app_id: "..." # or YUANBAO_APP_ID app_secret: "..." # or YUANBAO_APP_SECRET bot_id: "..." # or YUANBAO_BOT_ID (optional, returned by sign-token) ws_url: "wss://..." # or YUANBAO_WS_URL api_domain: "https://..." # or YUANBAO_API_DOMAIN """ from __future__ import annotations import asyncio import base64 import binascii import collections import dataclasses import hashlib import hmac import json import logging import os import re import secrets import time import urllib.parse import uuid from datetime import datetime, timezone, timedelta from enum import Enum from pathlib import Path from abc import ABC, abstractmethod from typing import Any, Callable, ClassVar, Dict, Iterator, List, Optional, Tuple import sys import httpx try: import websockets import websockets.exceptions WEBSOCKETS_AVAILABLE = True except ImportError: WEBSOCKETS_AVAILABLE = False websockets = None # type: ignore[assignment] from gateway.config import Platform, PlatformConfig from gateway.platforms.base import ( BasePlatformAdapter, MessageEvent, MessageType, SendResult, cache_document_from_bytes, cache_image_from_bytes, cache_video_from_bytes, ) from gateway.platforms import helpers as _mdchunk from gateway.platforms.helpers import MessageDeduplicator from gateway.platforms.yuanbao_media import ( download_url as media_download_url, get_cos_credentials, upload_to_cos, build_image_msg_body, build_file_msg_body, guess_mime_type, md5_hex, ) from gateway.platforms.yuanbao_proto import ( CMD_TYPE, _fields_to_dict, _get_string, _get_varint, _parse_fields, WS_HEARTBEAT_RUNNING, WS_HEARTBEAT_FINISH, HERMES_INSTANCE_ID, decode_conn_msg, decode_inbound_push, decode_forward_msg_data, decode_query_group_info_rsp, decode_get_group_member_list_rsp, encode_auth_bind, encode_ping, encode_push_ack, encode_send_c2c_message, encode_send_group_message, encode_send_private_heartbeat, encode_send_group_heartbeat, encode_query_group_info, encode_get_group_member_list, next_seq_no, ) from gateway.session import build_session_key logger = logging.getLogger(__name__) # Version / platform constants (AUTH_BIND and sign-token headers) try: from hermes_cli import __version__ as _HERMES_VERSION except ImportError: _HERMES_VERSION = "0.0.0" _APP_VERSION = _HERMES_VERSION _BOT_VERSION = _HERMES_VERSION _YUANBAO_INSTANCE_ID = str(HERMES_INSTANCE_ID) _OPERATION_SYSTEM = sys.platform DEFAULT_WS_GATEWAY_URL = "wss://bot-wss.yuanbao.tencent.com/wss/connection" DEFAULT_API_DOMAIN = "https://bot.yuanbao.tencent.com" HEARTBEAT_INTERVAL_SECONDS = 30.0 CONNECT_TIMEOUT_SECONDS = 15.0 AUTH_TIMEOUT_SECONDS = 10.0 MAX_RECONNECT_ATTEMPTS = 100 DEFAULT_SEND_TIMEOUT = 30.0 # WS biz request timeout # Bound on the WS close handshake at teardown/reconnect: websockets' own close_timeout (5s) waits # for the server's close echo, which an idle server never sends, stalling shutdown. A responsive # server finishes well under 1s, so this only caps the pathological hang. WS_CLOSE_TIMEOUT_S = 1.0 # Close codes that indicate permanent errors — do NOT reconnect. NO_RECONNECT_CLOSE_CODES = {4012, 4013, 4014, 4018, 4019, 4021} # Heartbeat timeout threshold — N consecutive missed pongs trigger reconnect. HEARTBEAT_TIMEOUT_THRESHOLD = 2 # Reply Heartbeat configuration REPLY_HEARTBEAT_INTERVAL_S = 2.0 # Send RUNNING every 2 seconds REPLY_HEARTBEAT_TIMEOUT_S = 30.0 # Auto-stop after 30 seconds of inactivity # Slow-response hint: push a waiting message when agent produces no data for this duration (seconds) SLOW_RESPONSE_TIMEOUT_S = 120.0 SLOW_RESPONSE_MESSAGE = "任务有点复杂,正在努力处理中,请耐心等待..." # Resource anchors in transcript text: [image|ybres:abc] [file:report.pdf|ybres:xyz] [voice|ybres:…] _YB_RES_REF_RE = re.compile(r"\[(image|voice|video|file(?::[^|\]]*)?)\|ybres:([A-Za-z0-9_\-]+)\]") # Anchors after local download: [image: /path] [file: report.pdf → /path] [video: /path] _YB_LOCAL_MEDIA_RE = re.compile(r"\[(\w+):[^\]]*?(/[^\]]+?)\s*\]") # Media kinds that can be resolved and injected into the model context _RESOLVABLE_MEDIA_KINDS = frozenset({"image", "file", "video"}) def _iter_ybres_refs(matches) -> Iterator[Tuple[str, str, str]]: """Turn ``_YB_RES_REF_RE`` matches into ``(rid, kind, filename)`` for resolvable kinds only.""" for m in matches: kind, _, filename = m.group(1).partition(":") kind = kind.strip() if kind in _RESOLVABLE_MEDIA_KINDS: yield m.group(2), kind, filename.strip() # Strip page indicators like (1/3) appended by BasePlatformAdapter _INDICATOR_RE = re.compile(r'\s*\(\d+/\d+\)$') # Observed-media backfill: how many recent transcript messages to scan OBSERVED_MEDIA_BACKFILL_LOOKBACK = 50 # Max number of resource references to resolve per inbound turn OBSERVED_MEDIA_BACKFILL_MAX_RESOLVE_PER_TURN = 12 # Inbound media resolve concurrency (config: platforms.yuanbao.extra.media_resolve_concurrency). # 1 = sequential rollback knob; 6 = browser per-origin HTTP/1.1 ceiling; 12 = backfill cap. _DEFAULT_RESOLVE_CONCURRENCY = 6 _MIN_RESOLVE_CONCURRENCY = 1 _MAX_RESOLVE_CONCURRENCY = 12 class MarkdownProcessor: """Markdown chunking utilities — thin delegates to the shared fence-aware chunker in gateway.platforms.helpers; method names kept for existing call sites and tests.""" @staticmethod def has_unclosed_fence(text: str) -> bool: """Detect whether the text has unclosed code block fences.""" return _mdchunk.text_has_unclosed_fence(text) # -- Table detection --------------------------------------------------- @staticmethod def ends_with_table_row(text: str) -> bool: """Detect whether the text ends with a table row.""" return _mdchunk.text_ends_with_table_row(text) # -- Paragraph boundary splitting -------------------------------------- @staticmethod def split_at_paragraph_boundary( text: str, max_chars: int, len_fn: Optional[Callable[[str], int]] = None, ) -> tuple[str, str]: """Find the nearest paragraph boundary within max_chars; return (head, tail).""" return _mdchunk.split_at_paragraph_boundary(text, max_chars, len_fn=len_fn) # -- Core: chunk splitting --------------------------------------------- @classmethod def chunk_markdown_text( cls, text: str, max_chars: int = 4000, len_fn: Optional[Callable[[str], int]] = None, ) -> list[str]: """Split Markdown into <= max_chars chunks at paragraph boundaries, never inside a code fence or table (an oversized single block may exceed the limit).""" return _mdchunk.split_text_fence_aware(text, max_chars, len_fn, prefer_paragraphs=True, balance_fences=False) class SignManager: """Sign-token acquisition, caching, signing and retry. All state is class-level so one shared client serves the whole process.""" TOKEN_PATH = "/api/v5/robotLogic/sign-token" RETRYABLE_CODE = 10099 MAX_RETRIES = 3 RETRY_DELAY_S = 1.0 CACHE_REFRESH_MARGIN_S = 60 # treat as expiring this many seconds early HTTP_TIMEOUT_S = 10.0 _cache: dict[str, dict[str, Any]] = {} # app_key → {"token", "bot_id", "expire_ts", ...} # Per-app_key refresh locks, created lazily from async context so they bind to the running # loop; disconnect() clears them to avoid stale locks across reconnects. _locks: dict[str, asyncio.Lock] = {} @classmethod def get_refresh_lock(cls, app_key: str) -> asyncio.Lock: """Per-app_key refresh lock (create on demand). Call only from a running event loop.""" if app_key not in cls._locks: cls._locks[app_key] = asyncio.Lock() return cls._locks[app_key] @staticmethod def compute_signature(nonce: str, timestamp: str, app_key: str, app_secret: str) -> str: """HMAC-SHA256(key=app_secret, msg=nonce+timestamp+app_key+app_secret).hexdigest().""" plain = nonce + timestamp + app_key + app_secret return hmac.new(app_secret.encode(), plain.encode(), hashlib.sha256).hexdigest() @staticmethod def build_timestamp() -> str: """Beijing-time ISO-8601 timestamp without milliseconds (2006-01-02T15:04:05+08:00).""" bjtime = datetime.now(tz=timezone(timedelta(hours=8))) return bjtime.strftime("%Y-%m-%dT%H:%M:%S+08:00") @classmethod def is_cache_valid(cls, entry: dict[str, Any]) -> bool: """Determine whether the cache entry is valid (not expired with margin).""" return entry["expire_ts"] - time.time() > cls.CACHE_REFRESH_MARGIN_S @classmethod def clear_locks(cls) -> None: """Clear all per-app_key refresh locks (called on disconnect).""" cls._locks.clear() @classmethod def purge_expired(cls) -> int: """Drop expired token-cache entries (called lazily from get_token); returns count purged.""" now = time.time() expired_keys = [k for k, v in cls._cache.items() if now - v.get("expire_ts", 0) > 0] for k in expired_keys: cls._cache.pop(k, None) return len(expired_keys) @classmethod async def fetch( cls, app_key: str, app_secret: str, api_domain: str, route_env: str = "", ) -> dict[str, Any]: """Send sign-ticket HTTP request with auto-retry (up to MAX_RETRIES times).""" url = f"{api_domain.rstrip('/')}{cls.TOKEN_PATH}" async with httpx.AsyncClient(timeout=cls.HTTP_TIMEOUT_S) as client: for attempt in range(cls.MAX_RETRIES + 1): nonce = secrets.token_hex(16) timestamp = cls.build_timestamp() signature = cls.compute_signature(nonce, timestamp, app_key, app_secret) payload = { "app_key": app_key, "nonce": nonce, "signature": signature, "timestamp": timestamp, } headers = { "Content-Type": "application/json", "X-AppVersion": _APP_VERSION, "X-OperationSystem": _OPERATION_SYSTEM, "X-Instance-Id": _YUANBAO_INSTANCE_ID, "X-Bot-Version": _BOT_VERSION, } if route_env: headers["X-Route-Env"] = route_env logger.info( "Sign token request: url=%s%s", url, f" (retry {attempt}/{cls.MAX_RETRIES})" if attempt > 0 else "", ) response = await client.post(url, json=payload, headers=headers) if response.status_code != 200: body = response.text raise RuntimeError(f"Sign token API returned {response.status_code}: {body[:200]}") try: result_data: dict[str, Any] = response.json() except Exception as exc: raise ValueError(f"Sign token response parse error: {exc}") from exc code = result_data.get("code") if code == 0: data = result_data.get("data") if not isinstance(data, dict): raise ValueError(f"Sign token response missing 'data' field: {result_data}") logger.info("Sign token success: bot_id=%s", data.get("bot_id")) return data if code == cls.RETRYABLE_CODE and attempt < cls.MAX_RETRIES: logger.warning( "Sign token retryable: code=%s, retrying in %ss (attempt=%d/%d)", code, cls.RETRY_DELAY_S, attempt + 1, cls.MAX_RETRIES, ) await asyncio.sleep(cls.RETRY_DELAY_S) continue msg = result_data.get("msg", "") raise RuntimeError(f"Sign token error: code={code}, msg={msg}") raise RuntimeError("Sign token failed: max retries exceeded") @classmethod async def _fetch_into_cache(cls, app_key: str, app_secret: str, api_domain: str, route_env: str) -> None: data = await cls.fetch(app_key, app_secret, api_domain, route_env) duration: int = data.get("duration", 0) expire_ts = time.time() + duration if duration > 0 else time.time() + 3600 cls._cache[app_key] = { "token": data.get("token", ""), "bot_id": data.get("bot_id", ""), "duration": duration, "product": data.get("product", ""), "source": data.get("source", ""), "expire_ts": expire_ts, } @classmethod async def get_token( cls, app_key: str, app_secret: str, api_domain: str, route_env: str = "", ) -> dict[str, Any]: """Get WS auth token, served from cache while valid (with CACHE_REFRESH_MARGIN_S).""" cls.purge_expired() cached = cls._cache.get(app_key) if cached and cls.is_cache_valid(cached): remain = int(cached["expire_ts"] - time.time()) logger.info("Using cached token (%ds remaining)", remain) return dict(cached) async with cls.get_refresh_lock(app_key): cached = cls._cache.get(app_key) if cached and cls.is_cache_valid(cached): return dict(cached) await cls._fetch_into_cache(app_key, app_secret, api_domain, route_env) return dict(cls._cache[app_key]) @classmethod async def force_refresh( cls, app_key: str, app_secret: str, api_domain: str, route_env: str = "", ) -> dict[str, Any]: """Force refresh token (clear cache and re-sign).""" logger.warning("[force-refresh] Clearing cache and re-signing token: app_key=****%s", app_key[-4:]) async with cls.get_refresh_lock(app_key): cls._cache.pop(app_key, None) await cls._fetch_into_cache(app_key, app_secret, api_domain, route_env) return dict(cls._cache[app_key]) from dataclasses import dataclass, field as dc_field from gateway.platforms._shared import get_scoped_secret as _yb_secret @dataclass class InboundContext: """Mutable context passed through every inbound middleware in registration order.""" adapter: Any # YuanbaoAdapter (forward-ref avoids circular import) raw_frames: list = dc_field(default_factory=list) # Raw bytes frames (debounce-aggregated) # Populated by DecodeMiddleware push: Optional[dict] = None decoded_via: str = "" # "json" | "protobuf" # Extracted from push by FieldExtractMiddleware from_account: str = "" group_code: str = "" group_name: str = "" sender_nickname: str = "" msg_body: list = dc_field(default_factory=list) msg_id: str = "" cloud_custom_data: str = "" # Derived by ChatRoutingMiddleware chat_id: str = "" chat_type: str = "" # "dm" | "group" chat_name: str = "" # Populated by ContentExtractMiddleware raw_text: str = "" media_refs: list = dc_field(default_factory=list) # Populated by ExtractContentMiddleware for elem_type 1009 (WeChat forward). # Contains the parsed ForwardMsgData dict (sub_type / nick_name / msg list). forwarded_records: Optional[dict] = None # Owner command detection owner_command: Optional[str] = None # Source built by BuildSourceMiddleware source: Optional[Any] = None # SessionSource # Populated by ClassifyMessageTypeMiddleware msg_type: Optional[Any] = None # MessageType | YuanbaoMessageType # Populated by QuoteContextMiddleware reply_to_message_id: Optional[str] = None reply_to_text: Optional[str] = None quote_media_refs: list = dc_field(default_factory=list) # List of (rid, kind, filename) # Populated by MediaResolveMiddleware. Combined list of resolved local # paths from up to three sources (deduped, in this order): # 1) media carried by the current message (always), # 2) media from the quoted message (when reply_to_message_id is set), # 3) recent group-observed media (only when chat_type == "group" and no quote is present). media_urls: list = dc_field(default_factory=list) media_types: list = dc_field(default_factory=list) # Populated by GroupAttributionMiddleware channel_prompt: Optional[str] = None class InboundMiddleware(ABC): """Inbound middleware: set class-level ``name`` and implement ``handle(ctx, next_fn)``; ``await next_fn()`` continues the pipeline, returning without it stops.""" name: str = "" # Override in each subclass @abstractmethod async def handle(self, ctx: InboundContext, next_fn: Callable) -> None: """Process *ctx* and optionally call *next_fn* to continue the pipeline.""" async def __call__(self, ctx: InboundContext, next_fn: Callable) -> None: """Allow middleware instances to be called directly (duck-typing compat).""" return await self.handle(ctx, next_fn) def __repr__(self) -> str: return f"<{self.__class__.__name__} name={self.name!r}>" class InboundPipeline: """Onion-model middleware pipeline: named middlewares, ``when`` guards, use_before/use_after/ remove. Accepts ``InboundMiddleware`` instances or plain ``async def(ctx, next_fn)`` callables.""" def __init__(self) -> None: self._middlewares: list = [] # list of (name, handler, when_fn | None) @staticmethod def _normalize(name_or_mw, handler=None): """(InboundMiddleware,) or (name, handler) → (name, callable).""" if isinstance(name_or_mw, InboundMiddleware): return name_or_mw.name, name_or_mw return name_or_mw, handler def use(self, name_or_mw, handler=None, when=None) -> "InboundPipeline": """Append ``pipeline.use(SomeMiddleware())`` or ``pipeline.use("name", fn)``.""" name, h = self._normalize(name_or_mw, handler) self._middlewares.append((name, h, when)) return self def _insert_relative(self, target: str, offset: int, name_or_mw, handler, when) -> "InboundPipeline": name, h = self._normalize(name_or_mw, handler) idx = next((i for i, (n, _, _) in enumerate(self._middlewares) if n == target), None) entry = (name, h, when) if idx is None: self._middlewares.append(entry) else: self._middlewares.insert(idx + offset, entry) return self def use_before(self, target: str, name_or_mw, handler=None, when=None) -> "InboundPipeline": """Insert a middleware before *target* (by name). Appends if not found.""" return self._insert_relative(target, 0, name_or_mw, handler, when) def use_after(self, target: str, name_or_mw, handler=None, when=None) -> "InboundPipeline": """Insert a middleware after *target* (by name). Appends if not found.""" return self._insert_relative(target, 1, name_or_mw, handler, when) def remove(self, name: str) -> "InboundPipeline": """Remove a middleware by name.""" self._middlewares = [(n, h, w) for n, h, w in self._middlewares if n != name] return self @property def middleware_names(self) -> list: """Return ordered list of registered middleware names (for testing).""" return [n for n, _, _ in self._middlewares] async def execute(self, ctx: InboundContext) -> None: """Run all middlewares in order. Each middleware receives ``(ctx, next_fn)``.""" chain = self._middlewares index = 0 async def next_fn() -> None: nonlocal index while index < len(chain): name, handler, when_fn = chain[index] index += 1 if when_fn is not None and not when_fn(ctx): continue try: await handler(ctx, next_fn) except Exception: logger.error("[InboundPipeline] middleware [%s] error", name, exc_info=True) raise return await next_fn() class DecodeMiddleware(InboundMiddleware): """Decode raw inbound frames (JSON or protobuf via ``decode_inbound_push``) into ctx.push.""" name = "decode" @staticmethod def convert_json_msg_body(raw_body: list) -> list: """Normalize JSON msg_body (PascalCase or snake_case keys) to [{"msg_type", "msg_content"}].""" result = [] for item in raw_body or []: if not isinstance(item, dict): continue msg_type = item.get("msg_type") or item.get("MsgType", "") msg_content = item.get("msg_content") or item.get("MsgContent", {}) if isinstance(msg_content, str): try: msg_content = json.loads(msg_content) except Exception: msg_content = {"text": msg_content} result.append({"msg_type": msg_type, "msg_content": msg_content or {}}) return result @staticmethod def json_sender_fields(raw_json: dict) -> Tuple[str, str]: """(from_account, group_code) accepting both Tencent IM PascalCase and internal snake_case keys.""" from_account = raw_json.get("from_account", "") or raw_json.get("From_Account", "") group_code = raw_json.get("group_code", "") or raw_json.get("GroupId", "") or raw_json.get("group_id", "") return from_account, group_code @staticmethod def parse_json_push(raw_json: dict) -> dict | None: """JSON push → dict shaped like ``decode_inbound_push`` output; accepts both the callback format (callback_command/from_account/msg_body) and legacy keys (GroupId/MsgSeq/MsgKey/MsgBody).""" if not raw_json: return None from_account, group_code = DecodeMiddleware.json_sender_fields(raw_json) msg_body_raw = (raw_json.get("msg_body", []) or raw_json.get("MsgBody", [])) msg_body = DecodeMiddleware.convert_json_msg_body(msg_body_raw) # Recall callbacks may have neither from_account nor msg_body. if not from_account and not msg_body and not raw_json.get("callback_command"): return None return { "callback_command": raw_json.get("callback_command", ""), "from_account": from_account, "to_account": raw_json.get("to_account", "") or raw_json.get("To_Account", ""), "sender_nickname": raw_json.get("sender_nickname", "") or raw_json.get("nick_name", ""), "group_code": group_code, "group_name": raw_json.get("group_name", ""), "msg_seq": raw_json.get("msg_seq", 0) or raw_json.get("MsgSeq", 0), "msg_id": raw_json.get("msg_id", "") or raw_json.get("msg_key", "") or raw_json.get("MsgKey", ""), "msg_body": msg_body, "cloud_custom_data": raw_json.get("cloud_custom_data", "") or raw_json.get("CloudCustomData", ""), "bot_owner_id": raw_json.get("bot_owner_id", "") or raw_json.get("botOwnerId", ""), "recall_msg_seq_list": raw_json.get("recall_msg_seq_list") or None, "trace_id": (raw_json.get("log_ext") or {}).get("trace_id", "") if isinstance(raw_json.get("log_ext"), dict) else "", } def _decode_single(self, adapter, data: bytes) -> tuple: """Decode a single raw frame into (push_dict, decoded_via) or (None, '').""" try: conn_json = json.loads(data.decode("utf-8")) except Exception: conn_json = None if isinstance(conn_json, dict): push = self.parse_json_push(conn_json) if push: return push, "json" else: try: push = decode_inbound_push(data) except Exception: push = None if push: return push, "protobuf" return None, "" async def handle(self, ctx: InboundContext, next_fn) -> None: data_list = ctx.raw_frames if not data_list: return # Stop pipeline — nothing to decode merged_push = None decoded_via = "" for data in data_list: 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 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: _sep = {"msg_type": "TIMTextElem", "msg_content": {"text": "\n"}} merged_push["msg_body"] = merged_push.get("msg_body", []) + [_sep] + extra_body logger.info( "[%s] Merged %d extra msg_body elements from aggregated push", ctx.adapter.name, len(extra_body), ) if not merged_push: return # Stop pipeline ctx.push = merged_push ctx.decoded_via = decoded_via 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", [])], ) logger.debug("[%s] Push payload: %s", ctx.adapter.name, ctx.push) await next_fn() class ExtractFieldsMiddleware(InboundMiddleware): """Extract common fields from ctx.push into ctx attributes.""" name = "extract-fields" async def handle(self, ctx: InboundContext, next_fn) -> None: push = ctx.push ctx.from_account = push.get("from_account", "") ctx.group_code = push.get("group_code", "") ctx.group_name = push.get("group_name", "") ctx.sender_nickname = push.get("sender_nickname", "") ctx.msg_body = push.get("msg_body", []) ctx.msg_id = push.get("msg_id", "") ctx.cloud_custom_data = push.get("cloud_custom_data", "") await next_fn() class DedupMiddleware(InboundMiddleware): """Inbound message deduplication.""" name = "dedup" async def handle(self, ctx: InboundContext, next_fn) -> None: if ctx.msg_id and ctx.adapter._dedup.is_duplicate(ctx.msg_id): logger.debug("[%s] Duplicate message ignored: msg_id=%s", ctx.adapter.name, ctx.msg_id) return # Stop pipeline await next_fn() class RecallGuardMiddleware(InboundMiddleware): """Recall callbacks (Group.CallbackAfterRecallMsg / C2C.CallbackAfterMsgWithDraw). A: in transcript → redact; B: not in transcript → system note; C: being processed → interrupt + delayed redact.""" name = "recall_guard" _RECALL_COMMANDS = frozenset({"Group.CallbackAfterRecallMsg", "C2C.CallbackAfterMsgWithDraw"}) _REDACTED = "[This message was recalled/withdrawn by the sender; original content removed]" async def handle(self, ctx: InboundContext, next_fn) -> None: cmd = (ctx.push or {}).get("callback_command", "") if cmd not in self._RECALL_COMMANDS: await next_fn() return self._handle_recall(ctx, cmd) @staticmethod def _build_source(adapter, group_code: str, from_account: str): return adapter.build_source( chat_id=(f"group:{group_code}" if group_code else f"direct:{from_account}"), chat_type="group" if group_code else "dm", user_id=from_account or None, thread_id="main" if group_code else None, ) def _handle_recall(self, ctx: InboundContext, cmd: str) -> None: adapter = ctx.adapter push = ctx.push or {} 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") 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) return group_code = (push.get("group_code") or "").strip() from_account = (push.get("from_account") or "").strip() for seq_entry in seq_list: recalled_id = seq_entry.get("msg_id") or str(seq_entry.get("msg_seq") or "") if not recalled_id: continue matched_sk = self._find_processing_session(adapter, recalled_id) if matched_sk is not None: self._interrupt_for_recall(adapter, matched_sk, recalled_id, group_code, from_account) else: recalled_content = adapter._msg_content_cache.get(recalled_id) self._patch_transcript(adapter, recalled_id, group_code, from_account, recalled_content) # -- Branch C: interrupt currently-processing message --------------- @staticmethod def _find_processing_session(adapter, recalled_id: str) -> Optional[str]: for sk, mid in adapter._processing_msg_ids.items(): if mid == recalled_id and sk in adapter._active_sessions: return sk return None @classmethod def _interrupt_for_recall(cls, adapter, session_key: str, recalled_id: str, group_code: str, from_account: str) -> None: where = f"group {group_code}" if group_code else f"direct chat with {from_account}" recall_text = ( f"[CRITICAL — MESSAGE RECALLED] The user message that triggered " f"your current task (message_id=\"{recalled_id}\") in {where} has " f"been recalled/withdrawn by the sender. " f"IGNORE any prior system note asking you to finish processing " f"tool results — the original request is void. " f"Do NOT continue the task, do NOT call more tools, do NOT " f"reference the recalled content. " f"Reply only with a brief acknowledgment such as " f"\"The message has been recalled.\" in the " f"language the user was using." ) synth_event = MessageEvent( text=recall_text, message_type=MessageType.TEXT, source=cls._build_source(adapter, group_code, from_account), internal=True, ) # Set pending + signal directly (bypass handle_message to avoid busy-ack). # May overwrite a user message pending in the same ~200ms window — acceptable. adapter._pending_messages[session_key] = synth_event active_event = adapter._active_sessions.get(session_key) if active_event is not None: active_event.set() logger.info("[%s] Recall interrupt: msg_id=%s session=%s", adapter.name, recalled_id, session_key[:30]) # The interrupted turn will persist the recalled content *after* our # interrupt — schedule a delayed redaction to clean it up. recalled_text = adapter._processing_msg_texts.get(session_key, "") if recalled_text: cls._schedule_content_redact(adapter, session_key, recalled_text, group_code, from_account) @classmethod def _schedule_content_redact(cls, adapter, session_key: str, recalled_text: str, group_code: str, from_account: str) -> None: async def _redact() -> None: store = getattr(adapter, "_session_store", None) if not store: return try: sid = store.get_or_create_session( cls._build_source(adapter, group_code, from_account), ).session_id except Exception: return # Poll until the recalled content appears in transcript — the # interrupted turn hasn't finished writing yet when scheduled. for _ in range(30): await asyncio.sleep(0.5) try: transcript = store.load_transcript(sid) except Exception: continue for entry in transcript: if entry.get("role") == "user" and entry.get("content") == recalled_text: entry["content"] = cls._REDACTED try: store.rewrite_transcript(sid, transcript, active_only=True) logger.info("[%s] Recall redact: session %s", adapter.name, session_key[:30]) except Exception as exc: logger.warning("[%s] Recall redact failed: %s", adapter.name, exc) return logger.debug("[%s] Recall redact: content not found after polling, session %s", adapter.name, session_key[:30]) task = asyncio.create_task(_redact()) adapter._background_tasks.add(task) task.add_done_callback(adapter._background_tasks.discard) # -- Branch A/B: patch transcript (session idle) -------------------- @classmethod def _patch_transcript(cls, adapter, recalled_id: str, group_code: str, from_account: str, recalled_content: Optional[str] = None) -> None: store = getattr(adapter, "_session_store", None) if not store: return try: sid = store.get_or_create_session(cls._build_source(adapter, group_code, from_account)).session_id except Exception as exc: logger.warning("[%s] Recall: failed to resolve session: %s", adapter.name, exc) return try: transcript = store.load_transcript(sid) except Exception as exc: logger.warning("[%s] Recall: failed to load transcript: %s", adapter.name, exc) return # A1: exact platform message_id match (rows persisted with a platform_message_id). target = None branch_label = "" for entry in transcript: if entry.get("message_id") == recalled_id: target = entry branch_label = "branch A1: id match" break # A2: content-match fallback for rows without a platform id (agent-processed @bot # messages — run.py doesn't carry msg_id through — or pre-column rows). if target is None and recalled_content: for entry in transcript: if entry.get("role") == "user" and entry.get("content") == recalled_content: target = entry branch_label = "branch A2: content match" break if target is not None: target["content"] = cls._REDACTED try: store.rewrite_transcript(sid, transcript, active_only=True) logger.info("[%s] Recall: redacted msg_id=%s (%s)", adapter.name, recalled_id, branch_label) except Exception as exc: logger.warning("[%s] Recall: rewrite_transcript failed: %s", adapter.name, exc) return # Branch B: not found in transcript → append system note store.append_to_transcript(sid, { "role": "system", "content": f'[recall] message_id="{recalled_id}" has been recalled; do not quote or reference it.', "timestamp": datetime.now(tz=timezone.utc).isoformat(), }) logger.info("[%s] Recall: system note for msg_id=%s (branch B)", adapter.name, recalled_id) class SkipSelfMiddleware(InboundMiddleware): """Filter out bot's own messages.""" name = "skip-self" @staticmethod def _is_self_reference(from_account: str, bot_id: Optional[str]) -> bool: """Detect whether the message is from the bot itself.""" if not from_account or not bot_id: return False return from_account == bot_id async def handle(self, ctx: InboundContext, next_fn) -> None: if self._is_self_reference(ctx.from_account, ctx.adapter._bot_id): logger.debug("[%s] Ignoring self-sent message from %s", ctx.adapter.name, ctx.from_account) return # Stop pipeline await next_fn() class ChatRoutingMiddleware(InboundMiddleware): """Determine chat_id, chat_type, chat_name from push fields.""" name = "chat-routing" async def handle(self, ctx: InboundContext, next_fn) -> None: if ctx.group_code: ctx.chat_id = f"group:{ctx.group_code}" ctx.chat_type = "group" ctx.chat_name = ctx.group_name or ctx.group_code else: ctx.chat_id = f"direct:{ctx.from_account}" ctx.chat_type = "dm" ctx.chat_name = ctx.sender_nickname or ctx.from_account await next_fn() class AccessPolicy: """DM / group access rules shared by inbound middleware and outbound ``send_dm``.""" def __init__( self, dm_policy: str, dm_allow_from: list[str], group_policy: str, group_allow_from: list[str], ) -> None: self._dm_policy = dm_policy self._dm_allow_from = dm_allow_from self._group_policy = group_policy self._group_allow_from = group_allow_from def _open_dm_opted_in(self) -> bool: if (_yb_secret("GATEWAY_ALLOW_ALL_USERS", "") or "").lower() in {"true", "1", "yes"}: return True return (_yb_secret("YUANBAO_ALLOW_ALL_USERS", "") or "").lower() in {"true", "1", "yes"} def _evaluate(self, policy: str, allow_from: list[str], principal: str, *, pairing: bool) -> bool: """Shared allow/deny rule; *pairing* is the verdict for the "pairing" policy.""" if policy == "allowlist": return principal in allow_from if policy == "pairing": return pairing if policy == "open": return self._open_dm_opted_in() return False # "disabled" or unknown def is_dm_allowed(self, sender_id: str) -> bool: """Strict DM authorization — pairing does not imply access.""" return self._evaluate(self._dm_policy, self._dm_allow_from, sender_id.strip(), pairing=False) def is_dm_intake_allowed(self, sender_id: str) -> bool: """Whether a DM may reach gateway intake (pairing handshake path).""" principal = str(sender_id or "").strip() if not principal: return False return self._evaluate(self._dm_policy, self._dm_allow_from, principal, pairing=True) def is_group_allowed(self, group_code: str) -> bool: """Platform-level group chat inbound filter (open / allowlist / disabled).""" return self._evaluate(self._group_policy, self._group_allow_from, group_code.strip(), pairing=False) @property def dm_policy(self) -> str: return self._dm_policy @property def group_policy(self) -> str: return self._group_policy class AccessGuardMiddleware(InboundMiddleware): """Platform-level DM/Group access control filter.""" name = "access-guard" async def handle(self, ctx: InboundContext, next_fn) -> None: adapter = ctx.adapter policy: AccessPolicy = adapter._access_policy if ctx.chat_type == "dm": if not policy.is_dm_intake_allowed(ctx.from_account): logger.debug( "[%s] DM from %s blocked by dm_policy=%s", adapter.name, ctx.from_account, policy.dm_policy, ) return # Stop pipeline elif ctx.chat_type == "group" and not policy.is_group_allowed(ctx.group_code): logger.debug( "[%s] Group %s blocked by group_policy=%s", adapter.name, ctx.group_code, policy.group_policy, ) return # Stop pipeline await next_fn() class AutoSetHomeMiddleware(InboundMiddleware): """Silently designate the first inbound conversation as home channel (config.yaml + env); a group home is upgraded by the first DM. Runs after GroupAtGuard so unaddressed group traffic never claims it; only strictly-authorized senders (allowlist / open opt-in / pairing-approved) may — intake-only pairing forwards must not. """ name = "auto-sethome" async def handle(self, ctx: InboundContext, next_fn) -> None: adapter = ctx.adapter if not adapter._auto_sethome_done and adapter._sender_may_designate_home(ctx): _cur_home = os.getenv("YUANBAO_HOME_CHANNEL", "") _should_set = ( not _cur_home or (_cur_home.startswith("group:") and ctx.chat_type == "dm") ) if ctx.chat_type == "dm": adapter._auto_sethome_done = True # DM seen — no further upgrades needed if _should_set: try: from hermes_constants import get_hermes_home from hermes_cli.config import atomic_config_write, read_user_config_raw _home = get_hermes_home() config_path = _home / "config.yaml" # Raw read: merged defaults must not be persisted to the user's file. user_config: dict = read_user_config_raw(config_path) user_config["YUANBAO_HOME_CHANNEL"] = ctx.chat_id atomic_config_write(config_path, user_config) os.environ["YUANBAO_HOME_CHANNEL"] = str(ctx.chat_id) logger.info( "[%s] Auto-sethome: designated %s (%s) as Yuanbao home channel", adapter.name, ctx.chat_id, ctx.chat_name, ) except Exception as e: logger.warning("[%s] Auto-sethome failed: %s", adapter.name, e) await next_fn() def _iter_custom_elems(msg_body: list) -> Iterator[Tuple[Any, dict]]: """Yield ``(custom, content)`` for each TIMCustomElem whose ``data`` parses as JSON (any type).""" for elem in msg_body or []: if not isinstance(elem, dict) or elem.get("msg_type") != "TIMCustomElem": continue content = elem.get("msg_content", {}) or {} if not isinstance(content, dict): continue data_str = content.get("data", "") if not data_str: continue try: custom = json.loads(data_str) except (json.JSONDecodeError, TypeError): continue yield custom, content class ExtractContentMiddleware(InboundMiddleware): """Extract raw text and media refs from msg_body.""" name = "extract-content" _CARD_CONTENT_MAX_LENGTH = 1000 @staticmethod def _format_shared_link(custom: dict) -> str: """Format elem_type 1010 (share card) into bracket-placeholder text.""" title = custom.get("title", "") link = custom.get("link", "") header = f"[share_card: {title} | {link}]" if link else f"[share_card: {title}]" lines = [header] max_len = ExtractContentMiddleware._CARD_CONTENT_MAX_LENGTH for field in ("card_content", "wechat_des"): val = custom.get(field) if val and isinstance(val, str): preview = val[:max_len] + "...(truncated)" if len(val) > max_len else val lines.append(f"Preview: {preview}") break if link: lines.append("[visit link for full content]") return "\n".join(lines) @staticmethod def _format_link_understanding(custom: dict) -> Optional[str]: """Format elem_type 1007 (link understanding card) into bracket-placeholder text.""" content = custom.get("content") if not content: return None try: parsed = json.loads(content) link = parsed.get("link") if isinstance(parsed, dict) else None except (json.JSONDecodeError, TypeError): link = None if not link or not isinstance(link, str): return None return f"[link: {link} | visit link for full content]" @staticmethod def _parse_resource_id(url: str) -> str: """resourceId from a Yuanbao resource URL's query string, or "".""" if not url: return "" try: query = urllib.parse.parse_qs(urllib.parse.urlparse(url).query) ids = query.get("resourceId") or query.get("resourceid") or [] return str(ids[0]).strip() if ids else "" except Exception: return "" @staticmethod def _pick_image_url(content: dict) -> str: """URL of the medium image (index 1), falling back to index 0, else "".""" arr = content.get("image_info_array") if not isinstance(arr, list): arr = [] image_info = None if len(arr) > 1 and isinstance(arr[1], dict): image_info = arr[1] elif len(arr) > 0 and isinstance(arr[0], dict): image_info = arr[0] return str((image_info or {}).get("url") or "").strip() @classmethod def _extract_text(cls, msg_body: list) -> str: """Plain text from MsgBody: text elems verbatim; media as ``[kind|ybres:RID]`` / ``[kind]`` (file: ``[file:{name}|ybres:RID]``); TIMFaceElem ``[emoji: name]``; custom elems by elem_type. Parts are space-joined.""" parts: list[str] = [] for elem in msg_body: elem_type: str = elem.get("msg_type", "") content: dict = elem.get("msg_content", {}) if elem_type == "TIMTextElem": text = content.get("text", "") if text: parts.append(text) elif elem_type == "TIMImageElem": rid = cls._parse_resource_id(cls._pick_image_url(content)) parts.append(f"[image|ybres:{rid}]" if rid else "[image]") elif elem_type == "TIMFileElem": filename = content.get("file_name", content.get("fileName", content.get("filename", ""))) file_url = str(content.get("url") or "").strip() rid = cls._parse_resource_id(file_url) if rid: parts.append(f"[file:{filename}|ybres:{rid}]" if filename else f"[file|ybres:{rid}]") else: parts.append(f"[file: {filename}]" if filename else "[file]") elif elem_type in ("TIMSoundElem", "TIMVideoFileElem"): kind = "voice" if elem_type == "TIMSoundElem" else "video" rid = cls._parse_resource_id(str(content.get("url") or "").strip()) parts.append(f"[{kind}|ybres:{rid}]" if rid else f"[{kind}]") elif elem_type == "TIMCustomElem": parts.append(cls._custom_elem_text(content.get("data", ""))) elif elem_type == "TIMFaceElem": raw_data = content.get("data", "") face_name = "" if raw_data: try: face_data = json.loads(raw_data) face_name = (face_data.get("name") or "").strip() except (json.JSONDecodeError, TypeError, AttributeError): pass parts.append(f"[emoji: {face_name}]" if face_name else "[emoji]") elif elem_type: # Unknown element type — include type as placeholder parts.append(f"[{elem_type}]") return " ".join(parts) if parts else "" _UNSUPPORTED = "[unsupported message type]" @classmethod def _custom_elem_text(cls, data_val: str) -> str: """Text for a TIMCustomElem by elem_type: 1002 mention, 1010 share card, 1007 link card, 1009 forwarded chat-record summary; malformed JSON is passed through verbatim.""" if not data_val: return cls._UNSUPPORTED try: custom = json.loads(data_val) except (json.JSONDecodeError, TypeError): return data_val if not isinstance(custom, dict): return cls._UNSUPPORTED ctype = custom.get("elem_type") if ctype == 1002: return custom.get("text", "[mention]") if ctype == 1010: return cls._format_shared_link(custom) if ctype == 1007: return cls._format_link_understanding(custom) or cls._UNSUPPORTED if ctype == 1009: return custom.get("text", "[chat record]") return cls._UNSUPPORTED @staticmethod def _rewrite_slash_command(text: str) -> str: """Strip and convert a leading full-width slash (Chinese IME) to ASCII so commands match.""" text = text.strip() if text.startswith('\uff0f'): # Full-width slash text = '/' + text[1:] return text @staticmethod def _extract_inbound_media_refs(msg_body: list) -> List[Dict[str, str]]: """Inbound image/file refs: ``[{"kind": "image", "url": ...}, {"kind": "file", "url": ..., "name": ...}]``.""" refs: List[Dict[str, str]] = [] for elem in msg_body or []: if not isinstance(elem, dict): continue msg_type = elem.get("msg_type", "") content = elem.get("msg_content", {}) or {} if not isinstance(content, dict): continue if msg_type == "TIMImageElem": image_url = ExtractContentMiddleware._pick_image_url(content) if image_url: refs.append({"kind": "image", "url": image_url}) continue if msg_type == "TIMFileElem": file_url = str(content.get("url") or "").strip() file_name = ( str(content.get("file_name") or "").strip() or str(content.get("fileName") or "").strip() or str(content.get("filename") or "").strip() ) if file_url: ref: Dict[str, str] = {"kind": "file", "url": file_url} if file_name: ref["name"] = file_name refs.append(ref) return refs @staticmethod def _extract_forwarded_records(msg_body: list, user_id: str = "") -> Optional[dict]: """ForwardMsgData for elem_type 1009 (WeChat forward), or None. Payload lives in ``msg_content.ext_map`` (pb field 999) under keys ``wexin_forward_msg_[id]_[userid]``; values are base64 protobuf (NOT JSON). First entry decoding to ``sub_type == 1`` wins. """ for custom, content in _iter_custom_elems(msg_body): if not (isinstance(custom, dict) and custom.get("elem_type") == 1009): continue ext_map = content.get("ext_map") or {} if not isinstance(ext_map, dict) or not ext_map: return None def _parse_value(value): if not isinstance(value, str) or not value: return None try: pb = base64.b64decode(value) except (binascii.Error, ValueError): return None data = decode_forward_msg_data(pb) if isinstance(data, dict) and data.get("sub_type") == 1: return data return None for key, value in ext_map.items(): if not key.startswith("wexin_forward_msg_"): continue parsed = _parse_value(value) if parsed is not None: return parsed return None async def handle(self, ctx: InboundContext, next_fn) -> None: ctx.raw_text = self._rewrite_slash_command(self._extract_text(ctx.msg_body)) ctx.media_refs = self._extract_inbound_media_refs(ctx.msg_body) ctx.forwarded_records = self._extract_forwarded_records(ctx.msg_body, ctx.from_account) await next_fn() class PlaceholderFilterMiddleware(InboundMiddleware): """Skip pure placeholder messages (e.g. '[image]' with no media).""" name = "placeholder-filter" SKIPPABLE_PLACEHOLDERS: frozenset = frozenset({ "[image]", "[图片]", "[file]", "[文件]", "[video]", "[视频]", "[voice]", "[语音]", }) @classmethod def is_skippable_placeholder(cls, text: str, media_count: int = 0) -> bool: """Detect whether the message is a pure placeholder (should be skipped).""" if media_count > 0: return False stripped = text.strip() return stripped in cls.SKIPPABLE_PLACEHOLDERS async def handle(self, ctx: InboundContext, next_fn) -> None: if self.is_skippable_placeholder(ctx.raw_text, len(ctx.media_refs)): logger.debug("[%s] Skipping placeholder message: %r", ctx.adapter.name, ctx.raw_text) return # Stop pipeline await next_fn() class OwnerCommandMiddleware(InboundMiddleware): """Bot-owner slash commands in groups: allowlisted commands skip @Bot; non-owner attempts are rejected.""" name = "owner-command" ALLOWLIST: frozenset = frozenset({ "/new", "/reset", "/retry", "/undo", "/stop", "/approve", "/deny", "/bg", "/btw", "/queue", "/q", }) _rewrite_slash_command = staticmethod(ExtractContentMiddleware._rewrite_slash_command) @classmethod def _detect_owner_command( cls, *, push: dict, msg_body: list, chat_type: str, from_account: str, ) -> Tuple[Optional[str], Optional[str], bool]: """→ (cmd, cmd_line, is_owner); (None, None, False) when not an allowlisted command.""" if chat_type != "group" or not cls.ALLOWLIST: return None, None, False # Extract TIMTextElem: only do command recognition with exactly one text segment text_elems = [e for e in (msg_body or []) if e.get("msg_type") == "TIMTextElem"] if len(text_elems) != 1: return None, None, False text = (text_elems[0].get("msg_content") or {}).get("text", "") cmd_line = cls._rewrite_slash_command(text) if not cmd_line.startswith("/"): return None, None, False cmd = cmd_line.split(maxsplit=1)[0].lower() if cmd not in cls.ALLOWLIST: return None, None, False # Owner ⇔ push.from_account == push.bot_owner_id. These commands are privileged # (/approve, /stop, /reset…) — a non-owner must never be able to run them. owner_id = str((push or {}).get("bot_owner_id") or "").strip() is_owner = bool(owner_id) and owner_id == from_account return cmd, cmd_line, is_owner async def handle(self, ctx: InboundContext, next_fn) -> None: adapter = ctx.adapter matched_cmd, cmd_line, is_owner = self._detect_owner_command( push=ctx.push, msg_body=ctx.msg_body, chat_type=ctx.chat_type, from_account=ctx.from_account, ) if matched_cmd and not is_owner: logger.info( "[%s] Reject non-owner slash command: chat=%s from=%s cmd=%s", adapter.name, ctx.chat_id, ctx.from_account, matched_cmd, ) adapter._track_task(asyncio.create_task( adapter.send(ctx.chat_id, f"⚠️ {matched_cmd} is only available to the creator in private chat mode"), name=f"yuanbao-owner-cmd-denial-{matched_cmd}", )) return # Stop pipeline if matched_cmd and is_owner and cmd_line: logger.info( "[%s] Bot owner slash command: chat=%s from=%s cmd=%s", adapter.name, ctx.chat_id, ctx.from_account, matched_cmd, ) ctx.owner_command = matched_cmd ctx.raw_text = cmd_line # Override with clean command text await next_fn() class BuildSourceMiddleware(InboundMiddleware): """Build SessionSource from context fields.""" name = "build-source" async def handle(self, ctx: InboundContext, next_fn) -> None: adapter = ctx.adapter ctx.source = adapter.build_source( chat_id=ctx.chat_id, chat_type=ctx.chat_type, chat_name=ctx.chat_name, user_id=ctx.from_account or None, user_name=ctx.sender_nickname or ctx.from_account, thread_id="main" if ctx.chat_type == "group" else None, ) await next_fn() class GroupAtGuardMiddleware(InboundMiddleware): """Group chat: observe non-@bot messages into the transcript; only @Bot (or owner commands) proceed.""" name = "group-at-guard" @staticmethod def _iter_bot_mentions(msg_body: list, bot_id: Optional[str]) -> Iterator[dict]: """Yield @bot elems: TIMCustomElem whose data JSON has elem_type 1002 and user_id == bot_id.""" if not bot_id: return for custom, _content in _iter_custom_elems(msg_body): if custom.get("elem_type") == 1002 and custom.get("user_id") == bot_id: yield custom @classmethod def _is_at_bot(cls, msg_body: list, bot_id: Optional[str]) -> bool: return any(True for _ in cls._iter_bot_mentions(msg_body, bot_id)) @classmethod def _extract_bot_mention_text(cls, msg_body: list, bot_id: Optional[str]) -> str: """Display text used to @-mention this bot (e.g. ``@yuanbao-bot``), or "".""" for custom in cls._iter_bot_mentions(msg_body, bot_id): mention_text = str(custom.get("text") or "").strip() if mention_text: return mention_text return "" @staticmethod def _build_group_channel_prompt(msg_body: list, bot_id: Optional[str]) -> str: """Build a per-turn group-chat prompt that highlights which message to respond to.""" bid = str(bot_id or "unknown") bot_mention = GroupAtGuardMiddleware._extract_bot_mention_text(msg_body, bot_id) or "unknown" return ( "You are handling a Yuanbao group chat message.\n" f"- Your identity: user_id={bid}, @-mention name in this group={bot_mention}\n" "- Lines in history prefixed with `[nickname|user_id]` are observed group context " "and are not necessarily addressed to you.\n" "- Treat only the current new message as a request explicitly directed at you, " "and answer it directly." ) @classmethod def _observe_group_message( cls, adapter, source, sender_display: str, text: str, *, ctx: InboundContext, msg_id: Optional[str] = None, forwarded_records: Optional[dict] = None, ) -> None: """Record a group message as ``role: user`` ``[nickname|user_id]\\n`` without invoking the agent, so later @bot turns see the full conversation.""" store = getattr(adapter, "_session_store", None) if not store: return try: session_entry = store.get_or_create_session(source) user_id = source.user_id or "unknown" body_text = text if forwarded_records: summary = ForwardedRecordsParseMiddleware.build_forward_text( forwarded_records, ctx=ctx, is_dispatch=False, ) if summary: body_text = f"{text}\n{summary}" if text else summary attributed = f"[{sender_display}|{user_id}]\n{body_text}" entry: dict = { "role": "user", "content": attributed, "timestamp": datetime.now(tz=timezone.utc).isoformat(), "observed": True, } if msg_id: entry["message_id"] = msg_id store.append_to_transcript(session_entry.session_id, entry) except Exception as exc: logger.warning("[%s] Failed to observe group message: %s", adapter.name, exc) async def handle(self, ctx: InboundContext, next_fn) -> None: adapter = ctx.adapter if ctx.chat_type == "group" and not ctx.owner_command and not self._is_at_bot(ctx.msg_body, adapter._bot_id): self._observe_group_message( adapter, ctx.source, ctx.sender_nickname or ctx.from_account, ctx.raw_text, msg_id=ctx.msg_id or None, forwarded_records=ctx.forwarded_records, ctx=ctx, ) logger.info( "[%s] Group message observed (no @bot): chat=%s from=%s", adapter.name, ctx.chat_id, ctx.from_account, ) return # Stop pipeline — message observed but not dispatched await next_fn() class GroupAttributionMiddleware(InboundMiddleware): """Group @bot turns: build channel_prompt, rewrite raw_text to ``[nickname|user_id]\\n`` (matches observed-history format) and clear ``source.user_name`` to suppress the runner's ``[user_name]`` prefix.""" name = "group-attribution" async def handle(self, ctx: InboundContext, next_fn) -> None: if ctx.chat_type == "group" and not ctx.owner_command: adapter = ctx.adapter ctx.channel_prompt = GroupAtGuardMiddleware._build_group_channel_prompt( ctx.msg_body, adapter._bot_id, ) user_id_label = ctx.from_account or "unknown" nickname_label = ctx.sender_nickname or ctx.from_account or "unknown" ctx.raw_text = f"[{nickname_label}|{user_id_label}]\n{ctx.raw_text}" if ctx.source is not None: ctx.source = dataclasses.replace(ctx.source, user_name=None) await next_fn() class YuanbaoMessageType(Enum): """Yuanbao-local subtypes; coerced back to MessageType in DispatchMiddleware.""" CHAT_RECORD = "chat_record" _ELEM_MESSAGE_TYPES = { "TIMImageElem": MessageType.PHOTO, "TIMSoundElem": MessageType.VOICE, "TIMVideoFileElem": MessageType.VIDEO, "TIMFileElem": MessageType.DOCUMENT, } class ClassifyMessageTypeMiddleware(InboundMiddleware): """Determine MessageType from text content and msg_body elements.""" name = "classify-msg-type" @staticmethod def _classify(text: str, msg_body: list): """MessageType (or yuanbao-local YuanbaoMessageType) from text and msg_body.""" if text.startswith("/"): return MessageType.COMMAND for elem in msg_body: etype = elem.get("msg_type", "") mapped = _ELEM_MESSAGE_TYPES.get(etype) if mapped is not None: return mapped if etype == "TIMCustomElem": data_str = (elem.get("msg_content") or {}).get("data", "") try: custom = json.loads(data_str) except (json.JSONDecodeError, TypeError): custom = None if isinstance(custom, dict) and custom.get("elem_type") == 1009: return YuanbaoMessageType.CHAT_RECORD return MessageType.TEXT async def handle(self, ctx: InboundContext, next_fn) -> None: ctx.msg_type = self._classify(ctx.raw_text, ctx.msg_body) await next_fn() class QuoteContextMiddleware(InboundMiddleware): """Extract quote/reply context from cloud_custom_data.""" name = "quote-context" def _extract_quote_context(self, cloud_custom_data: str) -> Tuple[Optional[str], Optional[str]]: """(quote_id, quote_text) from cloud_custom_data → MessageEvent.reply_to_*.""" if not cloud_custom_data: return None, None try: parsed = json.loads(cloud_custom_data) except (json.JSONDecodeError, TypeError): return None, None quote = parsed.get("quote") if isinstance(parsed, dict) else None if not isinstance(quote, dict): return None, None quote_id = str(quote.get("id") or "").strip() or None desc = str(quote.get("desc") or "").strip() sender = str(quote.get("sender_nickname") or quote.get("sender_id") or "").strip() quote_text = (f"{sender}: {desc}" if sender else desc) if desc else None return quote_id, quote_text async def _extract_media_refs_from_transcript( self, ctx: InboundContext ) -> List[Tuple[str, str, str]]: """``(rid, kind, filename)`` for ybres anchors in the quoted transcript message; [] when there is no reply_to id, no store/source, or no resolvable anchors.""" if ctx.reply_to_message_id is None: return [] adapter = ctx.adapter media_refs: List[Tuple[str, str, str]] = [] try: store = getattr(adapter, "_session_store", None) if not store or ctx.source is None: return [] session_entry = store.get_or_create_session(ctx.source) history = store.load_transcript(session_entry.session_id) for msg in reversed(history or []): mid = msg.get("message_id", "") if not mid or mid != ctx.reply_to_message_id: continue _content = msg.get("content", "") if isinstance(_content, str) and "|ybres:" in _content: media_refs.extend(_iter_ybres_refs(_YB_RES_REF_RE.finditer(_content))) break except Exception as exc: logger.warning( "[%s] quote transcript lookup failed: %s", getattr(adapter, "name", "yuanbao"), exc, ) return media_refs async def handle(self, ctx: InboundContext, next_fn) -> None: ctx.reply_to_message_id, ctx.reply_to_text = self._extract_quote_context(ctx.cloud_custom_data) ctx.quote_media_refs = await self._extract_media_refs_from_transcript(ctx) await next_fn() class ForwardedRecordsParseMiddleware(InboundMiddleware): """Deep-parse WeChat forwarded chat records (elem_type 1009) on ``ctx.forwarded_records``: render media as ``[kind|ybres:RID]``, append refs to ``ctx.media_refs`` and rewrite raw_text. No run-time fallback for earlier forwards — GroupAtGuard already rendered summaries at observe time. On any failure raw_text is left untouched.""" name = "forwarded-records-parse" async def handle(self, ctx: InboundContext, next_fn) -> None: try: if ctx.forwarded_records: await self._send_loading_heartbeat(ctx) ctx.raw_text = self.build_forward_text(ctx.forwarded_records, ctx=ctx, is_dispatch=True) except Exception as exc: logger.warning( "[%s] forwarded-records deep parse failed: %s", getattr(ctx.adapter, "name", "yuanbao"), exc, ) await next_fn() @staticmethod async def _send_loading_heartbeat(ctx: InboundContext) -> None: """Best-effort RUNNING heartbeat so the user sees a loading bubble.""" try: await ctx.adapter._outbound.heartbeat.send_heartbeat_once( ctx.chat_id, WS_HEARTBEAT_RUNNING, ) except Exception: pass @classmethod def _media_marker( cls, media: dict, plain_text: str = "", ) -> Tuple[str, Optional[Dict[str, str]]]: """One ``multimedia`` entry → ``(marker, ref)``: ``[kind|ybres:RID]`` + media_refs dict when a RID/URL is usable, else a plain ``[kind] name`` marker and ``ref=None``.""" media_type = (media.get("type", "") or media.get("doc_type", "")).strip().lower() url = str(media.get("url") or "").strip() media_id = str(media.get("media_id") or "").strip() 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 = media_id 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} 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 return f"[file] {file_name}".rstrip(), None if media_type == "url": # link share (e.g. WeChat article) — keep URL for the agent link_title = file_name or str(media.get("title") or "") return f"[link] {link_title} {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 FORWARD_MSG_TEXT_MAX_CHARS = 1000 # per-record text cap; record count is NOT capped @classmethod def _walk_forward_msgs( cls, forward_data: dict, ) -> Iterator[Tuple[str, str, List[Dict[str, str]]]]: """Yield ``(sender, body, refs)`` per ``ForwardMsgData['msg']`` record; body capped at FORWARD_MSG_TEXT_MAX_CHARS. ``refs`` keeps textual order — PatchAnchorsMiddleware relies on it.""" for msg in (forward_data.get("msg") if isinstance(forward_data, dict) else None) or []: if not isinstance(msg, dict): continue sender = msg.get("sender", "") plain_text = msg.get("plainText", "") msg_contents = msg.get("msgContent", []) or [] refs: List[Dict[str, str]] = [] if not msg_contents: rendered = plain_text else: parts: List[str] = [] for mc in msg_contents: if not isinstance(mc, dict): continue mc_type = mc.get("type", 0) # EnumMsgContentType: 1 TEXT, 2 MULTIMEDIA, 3 nested FORWARD if mc_type == 1: parts.append(mc.get("text", "")) elif mc_type == 2: for media in mc.get("multimedia", []) or []: if isinstance(media, dict): marker, ref = cls._media_marker(media, plain_text) parts.append(marker) if ref is not None: refs.append(ref) elif mc_type == 3: parts.append("[嵌套聊天记录]") else: if plain_text: parts.append(plain_text) rendered = " ".join(p for p in parts if p) or plain_text if len(rendered) > cls.FORWARD_MSG_TEXT_MAX_CHARS: rendered = rendered[: cls.FORWARD_MSG_TEXT_MAX_CHARS] + "…(已截断)" yield sender, rendered, refs @classmethod def build_forward_text( cls, forward_data: dict, *, ctx: InboundContext, is_dispatch: bool, ) -> str: """Render ``ForwardMsgData`` as ``发送人:正文`` lines with media markers. When ``is_dispatch``, refs go to ``ctx.media_refs`` and a ``用户附言:`` footer is added (observe-time callers skip both).""" nickname = ctx.sender_nickname or "用户" lines = [f"当前用户的昵称为{nickname}", "以下为用户的聊天记录"] for sender, body, refs in cls._walk_forward_msgs(forward_data): lines.append(f"{sender}:{body}") if is_dispatch: ctx.media_refs.extend(refs) text = "\n".join(lines) if is_dispatch and ctx.raw_text.strip(): text += f"\n\n用户附言:{ctx.raw_text.strip()}" return text class MediaResolveMiddleware(InboundMiddleware): """Resolve inbound media references to downloadable URLs.""" name = "media-resolve" # --- Resource download cache (keyed by resourceId) --- # Avoids redundant downloads of the same resource within the TTL window. _resource_cache: ClassVar[Dict[str, Tuple[str, str, float]]] = {} # rid -> (local_path, mime, ts) _RESOURCE_CACHE_TTL_S: ClassVar[int] = 24 * 60 * 60 # 24 hours _RESOURCE_CACHE_MAX_SIZE: ClassVar[int] = 256 @classmethod def _get_cached_resource(cls, resource_id: str) -> Optional[Tuple[str, str]]: """Return cached ``(local_path, mime)`` if still valid and file exists, else None.""" if not resource_id: return None entry = cls._resource_cache.get(resource_id) if entry is None: return None local_path, mime, ts = entry if time.time() - ts > cls._RESOURCE_CACHE_TTL_S: cls._resource_cache.pop(resource_id, None) return None # Verify the cached file still exists on disk (cache dir may be swept). if not os.path.isfile(local_path): cls._resource_cache.pop(resource_id, None) return None return local_path, mime @classmethod def _put_cached_resource(cls, resource_id: str, local_path: str, mime: str) -> None: """Store download result in cache. Evicts oldest entries when over capacity.""" if not resource_id: return if len(cls._resource_cache) >= cls._RESOURCE_CACHE_MAX_SIZE: # Drop the oldest 25% of entries by timestamp. sorted_keys = sorted(cls._resource_cache, key=lambda k: cls._resource_cache[k][2]) for k in sorted_keys[: cls._RESOURCE_CACHE_MAX_SIZE // 4]: cls._resource_cache.pop(k, None) cls._resource_cache[resource_id] = (local_path, mime, time.time()) @classmethod def _append_cached_resource( cls, adapter, resource_id: str, media_paths: List[str], mimes: List[str], ) -> bool: """Append a cached resource to output lists when available.""" hit = cls._get_cached_resource(resource_id) if hit is None: return False local_path, mime = hit logger.debug( "[%s] resource cache hit: rid=%s path=%s", adapter.name, resource_id, local_path, ) media_paths.append(local_path) mimes.append(mime) return True @staticmethod def _guess_image_ext_from_url(url: str) -> str: """Guess image extension from URL path.""" path = urllib.parse.urlparse(url).path ext = os.path.splitext(path)[1].lower() if ext in {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".heic", ".tiff"}: return ext return ".jpg" @staticmethod async def _fetch_resource_url(adapter, resource_id: str) -> str: """Low-level helper: exchange a ``resourceId`` for a direct download URL. Handles token retrieval, the ``/api/resource/v1/download`` API call, and a single 401-retry with token force-refresh. Raises on failure. """ resource_id = resource_id.strip() if not resource_id: raise RuntimeError("missing resource_id") token_data = await adapter._get_cached_token() token = str(token_data.get("token") or "").strip() source = str(token_data.get("source") or "web").strip() or "web" bot_id = str(token_data.get("bot_id") or adapter._bot_id or adapter._app_key).strip() if not token or not bot_id: raise RuntimeError("missing token or bot_id for resource download") api_url = f"{adapter._api_domain}/api/resource/v1/download" headers = { "Content-Type": "application/json", "X-ID": bot_id, "X-Token": token, "X-Source": source, } async with httpx.AsyncClient(timeout=15.0, follow_redirects=True) as client: for attempt in range(2): resp = await client.get(api_url, params={"resourceId": resource_id}, headers=headers) if resp.status_code == 401 and attempt == 0: # Force refresh token once on expiry and retry token_data = await SignManager.force_refresh( adapter._app_key, adapter._app_secret, adapter._api_domain, ) token = str(token_data.get("token") or "").strip() source = str(token_data.get("source") or source or "web").strip() or "web" bot_id = str(token_data.get("bot_id") or adapter._bot_id or adapter._app_key).strip() if not token or not bot_id: break headers["X-ID"] = bot_id headers["X-Token"] = token headers["X-Source"] = source continue resp.raise_for_status() payload = resp.json() code = payload.get("code") if code not in {None, 0}: raise RuntimeError( f"resource/v1/download failed: code={code}, msg={payload.get('msg', '')}" ) data = payload.get("data") if isinstance(payload.get("data"), dict) else payload real_url = str((data or {}).get("url") or (data or {}).get("realUrl") or "").strip() if real_url: return real_url raise RuntimeError("resource/v1/download missing url/realUrl") raise RuntimeError("resource/v1/download did not return a URL") @staticmethod async def _resolve_download_url(adapter, url: str) -> str: """Resolve Yuanbao resource placeholder to a directly fetchable real URL. Common URL patterns: https://hunyuan.tencent.com/api/resource/download?resourceId=... Direct GET returns 401; need business API: GET /api/resource/v1/download?resourceId=... """ try: parsed = urllib.parse.urlparse(url) except Exception: return url query = urllib.parse.parse_qs(parsed.query) resource_ids = query.get("resourceId") or query.get("resourceid") or [] resource_id = str(resource_ids[0]).strip() if resource_ids else "" if not resource_id: return url try: return await MediaResolveMiddleware._fetch_resource_url(adapter, resource_id) except Exception: return url @classmethod async def _download_and_cache( cls, adapter, *, fetch_url: str, kind: str, file_name: Optional[str] = None, log_tag: str = "", resource_id: str = "", ) -> Optional[Tuple[str, str]]: """Download a Yuanbao resource and cache locally. Returns ``(local_path, mime)`` or ``None``. When *resource_id* is provided, an in-memory cache keyed by resourceId is consulted first to skip redundant downloads of the same resource within the TTL window. """ if resource_id: hit = cls._get_cached_resource(resource_id) if hit is not None: logger.debug( "[%s] resource cache hit: rid=%s path=%s", adapter.name, resource_id, hit[0], ) return hit try: file_bytes, content_type = await media_download_url( fetch_url, max_size_mb=adapter.MEDIA_MAX_SIZE_MB, ) except Exception as exc: logger.warning( "[%s] inbound media download failed: kind=%s %s err=%s", adapter.name, kind, log_tag, exc, ) return None if kind == "image": ext = cls._guess_image_ext_from_url(fetch_url) try: local_path = cache_image_from_bytes(file_bytes, ext=ext) except ValueError as exc: logger.warning( "[%s] inbound image cache rejected: %s err=%s", adapter.name, log_tag, exc, ) return None mime = guess_mime_type(f"image{ext}") if not mime.startswith("image/"): mime = content_type if content_type.startswith("image/") else "image/jpeg" cls._put_cached_resource(resource_id, local_path, mime) return local_path, mime if kind == "video": # Yuanbao video resources carry no reliable extension; default to mp4. local_path = cache_video_from_bytes(file_bytes) mime = guess_mime_type(local_path) or ( content_type if content_type.startswith("video/") else "video/mp4" ) cls._put_cached_resource(resource_id, local_path, mime) return local_path, mime # kind == "file" if not file_name: parsed = urllib.parse.urlparse(fetch_url) file_name = os.path.basename(parsed.path) or "file" try: local_path = cache_document_from_bytes(file_bytes, file_name) except Exception as exc: logger.warning("[%s] inbound file cache failed: %s err=%s", adapter.name, log_tag, exc) return None mime = guess_mime_type(file_name) or content_type or "application/octet-stream" cls._put_cached_resource(resource_id, local_path, mime) return local_path, mime @classmethod async def _resolve_media_urls( cls, adapter, media_refs: List[Dict[str, str]] ) -> Tuple[List[str], List[str]]: """Resolve inbound media refs: download to local cache, return (local_paths, mime_types). Yuanbao COS hostnames resolve to private IPs, tripping the SSRF guard in vision_tools. We download ourselves and return local cache paths. Resolution runs with bounded concurrency (``adapter.media_resolve_concurrency``); see :meth:`_resolve_ybres_refs` for the same order-preserving / exception-isolated contract. """ # Pre-filter resolvable refs, preserving input order. media_urls: List[str] = [] media_types: List[str] = [] active: List[Tuple[str, str, str, str]] = [] for ref in media_refs: kind = str(ref.get("kind") or "").strip().lower() url = str(ref.get("url") or "").strip() filename = str(ref.get("name") or "").strip() if kind not in _RESOLVABLE_MEDIA_KINDS or not url: continue rid = ExtractContentMiddleware._parse_resource_id(url) if rid and cls._append_cached_resource(adapter, rid, media_urls, media_types): continue active.append((kind, url, filename, rid or "")) if not active: return media_urls, media_types async def _resolve_one(kind: str, url: str, filename: str, rid: str) -> Optional[Tuple[str, str]]: try: fetch_url = await cls._resolve_download_url(adapter, url) except Exception as exc: logger.warning( "[%s] inbound media resolve failed: kind=%s url=%s err=%s", adapter.name, kind, url, exc, ) return None return await cls._download_and_cache( adapter, fetch_url=fetch_url, kind=kind, file_name=filename or None, log_tag=f"placeholder_url={url[:80]}", resource_id=rid, ) def _crash_msg(item, result): kind, url, _filename, _rid = item return "[%s] inbound media resolve crashed: kind=%s url=%s err=%s", (adapter.name, kind, url[:80], result) await cls._gather_resolve(adapter, active, _resolve_one, _crash_msg, "media", media_urls, media_types) return media_urls, media_types @staticmethod async def _gather_resolve(adapter, active, resolve_one, crash_msg, scope, out_paths, out_mimes) -> None: """Run *resolve_one(*item)* over *active* under bounded concurrency; append successes in input order. ``return_exceptions=True`` isolates per-item failures; a batch summary log line with stable fields (concurrency vs elapsed_ms) is emitted for offline aggregation. """ semaphore = asyncio.Semaphore(adapter.media_resolve_concurrency) async def _guarded(item): async with semaphore: return await resolve_one(*item) _t0 = time.monotonic() results = await asyncio.gather(*(_guarded(item) for item in active), return_exceptions=True) _elapsed_ms = int((time.monotonic() - _t0) * 1000) _failed = 0 for item, result in zip(active, results): if isinstance(result, BaseException): fmt, args = crash_msg(item, result) logger.warning(fmt, *args) _failed += 1 continue if result is None: _failed += 1 continue path, mime = result out_paths.append(path) out_mimes.append(mime) logger.info( "[%s] media resolve batch: scope=%s concurrency=%d total=%d ok=%d failed=%d elapsed_ms=%d", adapter.name, scope, adapter.media_resolve_concurrency, len(active), len(out_paths), _failed, _elapsed_ms, ) @classmethod async def _resolve_ybres_refs( cls, adapter, refs: List[Tuple[str, str, str]], *, log_prefix: str, ) -> Tuple[List[str], List[str]]: """Resolve ``(rid, kind, filename)`` ybres tuples to local paths (bounded concurrency, input order preserved, per-rid failures isolated). Cache hits are served without a fetch.""" media_paths: List[str] = [] mimes: List[str] = [] active: List[Tuple[str, str, str]] = [] for rid, kind, filename in refs: if kind not in _RESOLVABLE_MEDIA_KINDS: continue if cls._append_cached_resource(adapter, rid, media_paths, mimes): continue active.append((rid, kind, filename)) if not active: return media_paths, mimes async def _resolve_one(rid: str, kind: str, filename: str) -> Optional[Tuple[str, str]]: try: fresh_url = await cls._fetch_resource_url(adapter, rid) except Exception as exc: logger.warning( "[%s] %s resolve failed: rid=%s kind=%s err=%s", adapter.name, log_prefix, rid, kind, exc, ) return None return await cls._download_and_cache( adapter, fetch_url=fresh_url, kind=kind, file_name=filename or None, log_tag=f"{log_prefix} rid={rid}", resource_id=rid, ) def _crash_msg(item, result): rid, kind, _filename = item return "[%s] %s resolve crashed: rid=%s kind=%s err=%s", (adapter.name, log_prefix, rid, kind, result) await cls._gather_resolve(adapter, active, _resolve_one, _crash_msg, "ybres", media_paths, mimes) return media_paths, mimes @classmethod async def _collect_observed_media(cls, adapter, source) -> Tuple[List[str], List[str]]: """Resolve recent observed image/file anchors from transcript into ``(local_paths, mimes)``.""" store = getattr(adapter, "_session_store", None) if not store: return [], [] try: session_entry = store.get_or_create_session(source) history = store.load_transcript(session_entry.session_id) except Exception as exc: logger.warning("[%s] Observed-media hydration setup failed: %s", adapter.name, exc) return [], [] if not history: return [], [] # Walk newest→oldest (matches within a message too) so the per-turn cap keeps the # *latest* refs; ``order`` is reversed back to chronological before resolving. window = history[-OBSERVED_MEDIA_BACKFILL_LOOKBACK:] order: List[Tuple[str, str, str]] = [] # (rid, kind, filename) seen: set = set() for msg in reversed(window): content = msg.get("content") if not isinstance(content, str) or "|ybres:" not in content: continue for rid, kind, filename in _iter_ybres_refs(reversed(list(_YB_RES_REF_RE.finditer(content)))): if rid in seen: continue seen.add(rid) order.append((rid, kind, filename)) if len(order) >= OBSERVED_MEDIA_BACKFILL_MAX_RESOLVE_PER_TURN: break if len(order) >= OBSERVED_MEDIA_BACKFILL_MAX_RESOLVE_PER_TURN: break order.reverse() if not order: return [], [] return await cls._resolve_ybres_refs(adapter, order, log_prefix="observed-media") @classmethod async def _resolve_quote_media( cls, adapter, quote_media_refs: List[Tuple[str, str, str]], ) -> Tuple[List[str], List[str]]: """Resolve ``(rid, kind, filename)`` anchors of the quoted message (from QuoteContextMiddleware).""" return await cls._resolve_ybres_refs(adapter, quote_media_refs, log_prefix="quote") @staticmethod def _collect_quote_local_media(ctx: InboundContext) -> Tuple[List[str], List[str]]: """DM quote fallback: ``(local_paths, mimes)`` for media PatchAnchorsMiddleware already rewrote to ``[image: /path]`` / ``[file: name → /path]`` on the original turn. Unresolved anchors were that turn's failure — no re-download here.""" paths: List[str] = [] mimes: List[str] = [] rid_key = ctx.reply_to_message_id if not rid_key: return paths, mimes cache = getattr(ctx.adapter, "_msg_content_cache", None) if not cache: return paths, mimes text = cache.get(rid_key) if not isinstance(text, str) or not text: return paths, mimes seen: set = set() for m in _YB_LOCAL_MEDIA_RE.finditer(text): kind = (m.group(1) or "").strip().lower() path = (m.group(2) or "").strip() if not path or path in seen: continue if not os.path.exists(path): continue seen.add(path) mime = guess_mime_type(os.path.basename(path)) or ( "image/jpeg" if kind == "image" else "application/octet-stream" ) paths.append(path) mimes.append(mime) return paths, mimes async def handle(self, ctx: InboundContext, next_fn) -> None: # In groups only @bot / owner-command turns reach here (GroupAtGuard short-circuits the # rest), so media download and observed-media hydration need no @bot re-check. adapter = ctx.adapter urls: List[str] = [] types: List[str] = [] seen: set = set() def _add_unique_pairs(pair_lists: Tuple[List[str], List[str]]) -> None: u_list, m_list = pair_lists for u, m in zip(u_list, m_list): if not u or u in seen: continue seen.add(u) urls.append(u) types.append(m) # 1) Media carried by the current message itself. own_pairs = await self._resolve_media_urls(adapter, ctx.media_refs) own_count = sum(1 for u in own_pairs[0] if u) _add_unique_pairs(own_pairs) # 2) Quoted media takes priority; else observed-media backfill in groups only (DM media # was already resolved on its own turn). if ctx.reply_to_message_id is not None: if ctx.quote_media_refs: _add_unique_pairs(await self._resolve_quote_media(adapter, ctx.quote_media_refs)) else: # DM rows carry no platform message_id → recover already-local media from the msg cache. _add_unique_pairs(self._collect_quote_local_media(ctx)) elif ctx.chat_type == "group": try: _add_unique_pairs(await self._collect_observed_media(adapter, ctx.source)) except Exception as exc: logger.warning( "[%s] observed-image hydration raised, continuing anyway: %s", adapter.name, exc, ) ctx.media_urls = urls ctx.media_types = types # Re-check placeholder using ``own_count``: placeholder text with only quote/observed # media (no fresh attachment of its own) is still skippable. if PlaceholderFilterMiddleware.is_skippable_placeholder(ctx.raw_text, own_count): logger.debug("[%s] Skip placeholder after media download: %r", adapter.name, ctx.raw_text) return # Stop pipeline await next_fn() class PatchAnchorsMiddleware(InboundMiddleware): """Replace ``[kind|ybres:RID]`` anchors in raw_text with the local paths MediaResolveMiddleware produced, so the transcript records usable paths. Only resolved media (paths starting with ``/``) are substituted; other anchors stay untouched.""" name = "patch-anchors" @staticmethod def _patch(text: str, urls: List[str], types: List[str]) -> str: if not text or not urls: return text patched = text for u, m in zip(urls, types): if not u.startswith("/"): continue anchor_match = _YB_RES_REF_RE.search(patched) if not anchor_match: break head = anchor_match.group(1) kind, _, filename = head.partition(":") kind = kind.strip() if kind == "image" and m.startswith("image/"): replacement = f"[image: {u}]" elif kind == "file": label = filename.strip() or os.path.basename(u) replacement = f"[file: {label} → {u}]" elif kind == "video": replacement = f"[video: {u}]" else: continue patched = patched[: anchor_match.start()] + replacement + patched[anchor_match.end():] return patched async def handle(self, ctx: InboundContext, next_fn) -> None: ctx.raw_text = self._patch(ctx.raw_text, ctx.media_urls, ctx.media_types) await next_fn() class DispatchMiddleware(InboundMiddleware): """Build MessageEvent and dispatch to AI handler.""" name = "dispatch" async def handle(self, ctx: InboundContext, next_fn) -> None: adapter = ctx.adapter _sk = build_session_key( ctx.source, group_sessions_per_user=adapter.config.extra.get("group_sessions_per_user", True), thread_sessions_per_user=adapter.config.extra.get("thread_sessions_per_user", False), ) async def _dispatch_inbound_event() -> None: event = MessageEvent( text=ctx.raw_text, message_type=( MessageType.DOCUMENT if any(mt.startswith(("application/", "text/")) for mt in ctx.media_types) # yuanbao-local subtypes (CHAT_RECORD) are deep-parsed into text → TEXT downstream else ctx.msg_type if isinstance(ctx.msg_type, MessageType) else MessageType.TEXT ), source=ctx.source, message_id=ctx.msg_id or None, raw_message=ctx.push, media_urls=list(ctx.media_urls), media_types=list(ctx.media_types), reply_to_message_id=ctx.reply_to_message_id, reply_to_text=ctx.reply_to_text, channel_prompt=ctx.channel_prompt, ) if _sk and ctx.msg_id: adapter._processing_msg_ids[_sk] = ctx.msg_id adapter._processing_msg_texts[_sk] = ctx.raw_text or "" if ctx.msg_id and ctx.raw_text: cache = adapter._msg_content_cache cache[ctx.msg_id] = ctx.raw_text if len(cache) > 200: for k in list(cache)[:len(cache) - 200]: del cache[k] await adapter.handle_message(event) if ctx.chat_type == "group": is_new = _sk not in adapter._group_queues queue = adapter._group_queues.setdefault(_sk, asyncio.Queue()) queue.put_nowait(_dispatch_inbound_event) logger.info( "[%s] Group message enqueued (qsize=%d) for %s", adapter.name, queue.qsize(), (_sk or "")[:50], ) if is_new: self._track_inbound(adapter, self._consume_group_queue(adapter, _sk), f"yuanbao-group-consumer-{(_sk or '')[:30]}") else: self._track_inbound(adapter, _dispatch_inbound_event(), f"yuanbao-inbound-{ctx.msg_id or 'unknown'}") await next_fn() @staticmethod def _track_inbound(adapter, coro, name: str) -> None: task = asyncio.create_task(coro, name=name) adapter._inbound_tasks.add(task) task.add_done_callback(adapter._inbound_tasks.discard) @staticmethod async def _consume_group_queue(adapter: "YuanbaoAdapter", session_key: str) -> None: """Drain the group queue one dispatch at a time, waiting for each to finish.""" _IDLE_TIMEOUT = 2.0 queue = adapter._group_queues.get(session_key) if not queue: return try: while True: try: dispatch_fn = await asyncio.wait_for(queue.get(), timeout=_IDLE_TIMEOUT) except asyncio.TimeoutError: break logger.debug( "[%s] Group queue: dispatching for %s (remaining=%d)", adapter.name, (session_key or "")[:50], queue.qsize(), ) try: await dispatch_fn() while session_key in adapter._active_sessions: await asyncio.sleep(0.1) except Exception: logger.exception("[%s] Group queue consumer error", adapter.name) finally: adapter._group_queues.pop(session_key, None) class InboundPipelineBuilder: """Assembles the default Yuanbao inbound pipeline (order matters).""" _DEFAULT_MIDDLEWARES: list[type] = [ DecodeMiddleware, ExtractFieldsMiddleware, RecallGuardMiddleware, DedupMiddleware, SkipSelfMiddleware, ChatRoutingMiddleware, AccessGuardMiddleware, ExtractContentMiddleware, PlaceholderFilterMiddleware, OwnerCommandMiddleware, BuildSourceMiddleware, GroupAtGuardMiddleware, AutoSetHomeMiddleware, GroupAttributionMiddleware, ClassifyMessageTypeMiddleware, QuoteContextMiddleware, ForwardedRecordsParseMiddleware, MediaResolveMiddleware, PatchAnchorsMiddleware, DispatchMiddleware, ] @classmethod def build(cls) -> InboundPipeline: """Build the default inbound message processing pipeline.""" pipeline = InboundPipeline() for mw_cls in cls._DEFAULT_MIDDLEWARES: pipeline.use(mw_cls()) return pipeline class ConnectionManager: """WebSocket lifecycle: open/close, AUTH_BIND, ping/pong heartbeat, receive loop, backoff reconnect.""" def __init__(self, adapter: "YuanbaoAdapter") -> None: self._adapter = adapter self._ws = None # websockets connection self._connect_id: Optional[str] = None self._heartbeat_task: Optional[asyncio.Task] = None self._recv_task: Optional[asyncio.Task] = None self._pending_acks: Dict[str, asyncio.Future] = {} self._pending_pong: Optional[asyncio.Future] = None self._consecutive_hb_timeouts: int = 0 self._reconnect_attempts: int = 0 self._reconnecting: bool = False # Debounce buffer aggregating multi-part inbound messages: sender key -> frames / timer self._inbound_buffer: Dict[str, list] = {} self._inbound_timers: Dict[str, asyncio.TimerHandle] = {} @property def ws(self): return self._ws @property def is_connected(self) -> bool: if self._ws is None: return False open_attr = getattr(self._ws, "open", None) if open_attr is True: return True if callable(open_attr): try: return bool(open_attr()) except Exception: return False return False async def open(self) -> bool: """sign-token → WS connect → AUTH_BIND → start loops. Returns True on success.""" adapter = self._adapter if not WEBSOCKETS_AVAILABLE: msg = "Yuanbao startup failed: 'websockets' package not installed" adapter._set_fatal_error("yuanbao_missing_dependency", msg, retryable=True) logger.warning("[%s] %s. Run: pip install websockets", adapter.name, msg) return False if not adapter._app_key or not adapter._app_secret: msg = "Yuanbao startup failed: YUANBAO_APP_ID and YUANBAO_APP_SECRET are required" adapter._set_fatal_error("yuanbao_missing_credentials", msg, retryable=False) logger.error("[%s] %s", adapter.name, msg) return False if self.is_connected: logger.debug("[%s] Already connected, skipping connect()", adapter.name) return True if not adapter._acquire_platform_lock('yuanbao-app-key', adapter._app_key, 'Yuanbao app key'): return False try: logger.info("[%s] Fetching sign token from %s", adapter.name, adapter._api_domain) token_data = await SignManager.get_token( adapter._app_key, adapter._app_secret, adapter._api_domain, route_env=adapter._route_env, ) self._apply_bot_id(token_data) logger.info("[%s] Connecting to %s", adapter.name, adapter._ws_url) if not await self._dial(token_data): return False self._reconnect_attempts = 0 adapter._mark_connected() adapter._loop = asyncio.get_running_loop() self._start_loops(cancel_existing=False) logger.info( "[%s] Connected. connectId=%s botId=%s", adapter.name, self._connect_id, adapter._bot_id, ) YuanbaoAdapter.set_active(adapter) return True except asyncio.TimeoutError: logger.error("[%s] Connection timed out", adapter.name) await self._cleanup_ws() adapter._release_platform_lock() return False except Exception as exc: logger.error("[%s] connect() failed: %s", adapter.name, exc, exc_info=True) await self._cleanup_ws() adapter._release_platform_lock() return False def _apply_bot_id(self, token_data: dict) -> None: """Adopt bot_id returned by the sign-token API, if any.""" if token_data.get("bot_id"): self._adapter._bot_id = str(token_data["bot_id"]) async def _dial(self, token_data: dict) -> bool: """Open the WS (built-in ping/pong disabled) and run AUTH_BIND; cleans up on auth failure.""" self._ws = await asyncio.wait_for( websockets.connect( # type: ignore[attr-defined] self._adapter._ws_url, ping_interval=None, ping_timeout=None, close_timeout=5, ), timeout=CONNECT_TIMEOUT_SECONDS, ) if not await self._authenticate(token_data): await self._cleanup_ws() return False return True _LOOPS = (("_heartbeat_task", "_heartbeat_loop", "heartbeat"), ("_recv_task", "_receive_loop", "recv")) def _start_loops(self, *, cancel_existing: bool) -> None: """(Re)start the heartbeat and receive loops for the current connect_id.""" for attr, coro_name, tag in self._LOOPS: old = getattr(self, attr) if cancel_existing and old and not old.done(): old.cancel() setattr(self, attr, asyncio.create_task( getattr(self, coro_name)(), name=f"yuanbao-{tag}-{self._connect_id}", )) async def close(self) -> None: """Cancel background tasks, fail pending futures, and close the WebSocket.""" for attr, _coro_name, _tag in self._LOOPS: task = getattr(self, attr) if task: task.cancel() try: await task except asyncio.CancelledError: pass setattr(self, attr, None) disc_exc = RuntimeError("YuanbaoAdapter disconnected") for fut in self._pending_acks.values(): if not fut.done(): fut.set_exception(disc_exc) self._pending_acks.clear() SignManager.clear_locks() # avoid stale locks bound to a previous event loop await self._cleanup_ws() async def _authenticate(self, token_data: dict) -> bool: """Send AUTH_BIND and read frames until BIND_ACK; False on failure/timeout.""" adapter = self._adapter if self._ws is None: return False token = token_data.get("token", "") uid = adapter._bot_id or token_data.get("bot_id", "") source = token_data.get("source") or "bot" route_env = adapter._route_env or token_data.get("route_env", "") or "" msg_id = str(uuid.uuid4()) auth_bytes = encode_auth_bind( biz_id="ybBot", uid=uid, source=source, token=token, msg_id=msg_id, app_version=_APP_VERSION, operation_system=_OPERATION_SYSTEM, bot_version=_BOT_VERSION, route_env=route_env, ) await self._ws.send(auth_bytes) 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 while True: remaining = deadline - _loop.time() if remaining <= 0: logger.error("[%s] AUTH_BIND timeout waiting for BIND_ACK", adapter.name) return False raw = await asyncio.wait_for(self._ws.recv(), timeout=remaining) if not isinstance(raw, (bytes, bytearray)): continue try: msg = decode_conn_msg(bytes(raw)) except Exception: continue head = msg.get("head", {}) cmd_type = head.get("cmd_type", -1) cmd = head.get("cmd", "") if cmd_type == CMD_TYPE["Response"] and cmd == "auth-bind": connect_id = self._extract_connect_id(msg) if connect_id: self._connect_id = connect_id logger.info("[%s] BIND_ACK received: connectId=%s", adapter.name, connect_id) return True else: logger.error("[%s] BIND_ACK missing connectId", adapter.name) return False except asyncio.TimeoutError: logger.error("[%s] AUTH_BIND timeout", adapter.name) return False except Exception as exc: logger.error("[%s] AUTH_BIND error: %s", adapter.name, exc, exc_info=True) return False def _pop_pending(self, msg_id: str) -> Optional[asyncio.Future]: """Pop the not-yet-done future registered for *msg_id*, if any.""" if msg_id and msg_id in self._pending_acks: fut = self._pending_acks.pop(msg_id) if not fut.done(): return fut return None def _extract_connect_id(self, decoded_msg: dict) -> Optional[str]: """Extract connectId from decoded BIND_ACK message.""" data: bytes = decoded_msg.get("data", b"") if not data: return None try: fdict = _fields_to_dict(_parse_fields(data)) code = _get_varint(fdict, 1) if code != 0: message = _get_string(fdict, 2) logger.error( "[%s] AuthBindRsp error: code=%d message=%r", self._adapter.name, code, message, ) return None connect_id = _get_string(fdict, 3) return connect_id if connect_id else None except Exception as exc: logger.warning("[%s] Failed to extract connectId: %s", self._adapter.name, exc) return None async def _heartbeat_loop(self) -> None: """Send HEARTBEAT (ping) every 30s; trigger reconnect after threshold misses.""" adapter = self._adapter try: while adapter._running: await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS) if self._ws is None: continue try: msg_id = str(uuid.uuid4()) ping_bytes = encode_ping(msg_id) loop = asyncio.get_running_loop() pong_future: asyncio.Future = loop.create_future() self._pending_pong = pong_future self._pending_acks[msg_id] = pong_future await self._ws.send(ping_bytes) logger.debug("[%s] PING sent (msg_id=%s)", adapter.name, msg_id) try: await asyncio.wait_for(pong_future, timeout=10.0) self._consecutive_hb_timeouts = 0 except asyncio.TimeoutError: self._pending_acks.pop(msg_id, None) self._consecutive_hb_timeouts += 1 logger.warning( "[%s] PONG timeout (%d/%d)", adapter.name, self._consecutive_hb_timeouts, HEARTBEAT_TIMEOUT_THRESHOLD, ) if self._consecutive_hb_timeouts >= HEARTBEAT_TIMEOUT_THRESHOLD: logger.warning("[%s] Heartbeat threshold exceeded, triggering reconnect", adapter.name) self.schedule_reconnect() return finally: self._pending_acks.pop(msg_id, None) self._pending_pong = None except Exception as exc: logger.debug("[%s] Heartbeat send failed: %s", adapter.name, exc) except asyncio.CancelledError: pass async def _receive_loop(self) -> None: """Read WS frames and dispatch by cmd_type.""" adapter = self._adapter try: async for raw in self._ws: # type: ignore[union-attr] if not isinstance(raw, (bytes, bytearray)): continue await self._handle_frame(bytes(raw)) except asyncio.CancelledError: pass except websockets.exceptions.ConnectionClosed as close_exc: # type: ignore[union-attr] close_code = getattr(close_exc, 'code', None) logger.warning( "[%s] WebSocket connection closed: code=%s reason=%s", adapter.name, close_code, getattr(close_exc, 'reason', ''), ) if close_code and close_code in NO_RECONNECT_CLOSE_CODES: logger.error( "[%s] Close code %d is non-recoverable, NOT reconnecting", adapter.name, close_code, ) adapter._mark_disconnected() else: self.schedule_reconnect() except Exception as exc: logger.warning("[%s] receive_loop exited: %s", adapter.name, exc) self.schedule_reconnect() async def _handle_frame(self, raw: bytes) -> None: """Handle a single WebSocket frame.""" adapter = self._adapter try: msg = decode_conn_msg(raw) except Exception as exc: logger.debug("[%s] Failed to decode frame: %s", adapter.name, exc) return head = msg.get("head", {}) cmd_type = head.get("cmd_type", -1) cmd = head.get("cmd", "") msg_id = head.get("msg_id", "") need_ack = head.get("need_ack", False) data: bytes = msg.get("data", b"") # HEARTBEAT_ACK if cmd_type == CMD_TYPE["Response"] and cmd == "ping": 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) return # Fire-and-forget heartbeat ACKs: nobody awaits them; discard to avoid "Unmatched" noise. if cmd_type == CMD_TYPE["Response"] and cmd in {"send_group_heartbeat", "send_private_heartbeat"}: logger.debug("[%s] Heartbeat ACK received: cmd=%s msg_id=%s", adapter.name, cmd, msg_id) return # Response to an outbound RPC call if cmd_type == CMD_TYPE["Response"]: if msg_id and msg_id in self._pending_acks: fut = self._pop_pending(msg_id) if fut is not None: result = {"head": head} if data: result["data"] = data fut.set_result(result) else: logger.debug("[%s] Unmatched Response: cmd=%s msg_id=%s", adapter.name, cmd, msg_id) return # Server-initiated Push if cmd_type == CMD_TYPE["Push"]: logger.info("[%s] Push received: cmd=%s msg_id=%s data_len=%d", adapter.name, cmd, msg_id, len(data)) if need_ack and self._ws is not None: try: ack_bytes = encode_push_ack(head) await self._ws.send(ack_bytes) except Exception as ack_exc: logger.debug("[%s] Failed to send PushAck: %s", adapter.name, ack_exc) if msg_id and msg_id in self._pending_acks: fut = self._pop_pending(msg_id) if fut is not None: try: fut.set_result(decode_inbound_push(data) if data else {"head": head}) except Exception as exc: fut.set_exception(exc) return if data: # genuine inbound message — dispatch to AI logger.info( "[%s] WS received inbound push, decoding and dispatching: cmd=%s, data_len=%d", adapter.name, cmd, len(data), ) self._push_to_inbound(data) return logger.debug( "[%s] Ignoring frame: cmd_type=%d cmd=%s msg_id=%s", adapter.name, cmd_type, cmd, msg_id, ) _DEBOUNCE_WINDOW: float = 1.5 # seconds to wait for companion messages def _extract_sender_key(self, raw_data: bytes) -> str: """Debounce key 'from_account:group_code' (JSON or protobuf), else a unique fallback.""" try: parsed = json.loads(raw_data.decode("utf-8")) if isinstance(parsed, dict): from_account, group_code = DecodeMiddleware.json_sender_fields(parsed) if from_account: return f"{from_account}:{group_code}" except Exception: pass try: push = decode_inbound_push(raw_data) if push: return f"{push.get('from_account', '')}:{push.get('group_code', '')}" except Exception: pass return f"__unknown_{id(raw_data)}" def _push_to_inbound(self, raw_data: bytes) -> None: """Debounced dispatch: frames from one sender within _DEBOUNCE_WINDOW run as ONE pipeline execution, merging multi-part messages (e.g. image + text pushed separately).""" key = self._extract_sender_key(raw_data) existing_timer = self._inbound_timers.pop(key, None) if existing_timer: existing_timer.cancel() self._inbound_buffer.setdefault(key, []).append(raw_data) logger.debug( "[%s] Debounce: buffered frame for key=%s, count=%d", self._adapter.name, key, len(self._inbound_buffer[key]), ) loop = asyncio.get_running_loop() self._inbound_timers[key] = loop.call_later(self._DEBOUNCE_WINDOW, self._flush_inbound_buffer, key) def _flush_inbound_buffer(self, key: str) -> None: """Flush the debounce buffer for a given key — execute the pipeline.""" self._inbound_timers.pop(key, None) data_list = self._inbound_buffer.pop(key, []) if not data_list: return adapter = self._adapter logger.info( "[%s] Debounce flush: key=%s, aggregated %d frames", adapter.name, key, len(data_list), ) ctx = InboundContext(adapter=adapter, raw_frames=data_list) adapter._track_task(asyncio.create_task( adapter._inbound_pipeline.execute(ctx), name=f"yuanbao-pipeline-{key}", )) async def send_biz_request( self, encoded_conn_msg: bytes, req_id: str, timeout: float = DEFAULT_SEND_TIMEOUT, ) -> dict: """Send a business request and await its response future (pending_acks[req_id]), cleaning up on exit.""" if self._ws is None: raise RuntimeError("Not connected") loop = asyncio.get_running_loop() future: asyncio.Future = loop.create_future() self._pending_acks[req_id] = future try: await self._ws.send(encoded_conn_msg) return await asyncio.wait_for(asyncio.shield(future), timeout=timeout) finally: self._pending_acks.pop(req_id, None) def schedule_reconnect(self) -> None: """Schedule a reconnect only if running and not already reconnecting.""" if self._adapter._running and not self._reconnecting: asyncio.create_task(self._reconnect_with_backoff()) async def _reconnect_with_backoff(self) -> bool: """Reconnect with exponential backoff (1s, 2s, 4s, … up to 60s).""" if self._reconnecting: logger.debug("[%s] Reconnect already in progress, skipping", self._adapter.name) return False self._reconnecting = True try: return await self._do_reconnect() finally: self._reconnecting = False async def _do_reconnect(self) -> bool: """Internal reconnect loop, called under the _reconnecting guard.""" adapter = self._adapter for attempt in range(MAX_RECONNECT_ATTEMPTS): self._reconnect_attempts = attempt + 1 wait = min(2 ** attempt, 60) logger.info( "[%s] Reconnect attempt %d/%d in %ds", adapter.name, attempt + 1, MAX_RECONNECT_ATTEMPTS, wait, ) await asyncio.sleep(wait) await self._cleanup_ws() try: token_data = await SignManager.force_refresh( adapter._app_key, adapter._app_secret, adapter._api_domain, route_env=adapter._route_env, ) self._apply_bot_id(token_data) if not await self._dial(token_data): logger.warning("[%s] Re-auth failed on attempt %d", adapter.name, attempt + 1) continue self._reconnect_attempts = 0 self._consecutive_hb_timeouts = 0 adapter._mark_connected() self._start_loops(cancel_existing=True) logger.info( "[%s] Reconnected on attempt %d. connectId=%s", adapter.name, attempt + 1, self._connect_id, ) YuanbaoAdapter.set_active(adapter) return True except asyncio.TimeoutError: logger.warning("[%s] Reconnect attempt %d timed out", adapter.name, attempt + 1) except Exception as exc: logger.warning( "[%s] Reconnect attempt %d failed: %s", adapter.name, attempt + 1, exc ) logger.error( "[%s] Giving up after %d reconnect attempts", adapter.name, MAX_RECONNECT_ATTEMPTS ) adapter._mark_disconnected() return False async def _cleanup_ws(self) -> None: """Close and clear the WS, bounded by WS_CLOSE_TIMEOUT_S so an unresponsive server can't stall teardown.""" ws = self._ws self._ws = None if ws is not None: try: await asyncio.wait_for(ws.close(), timeout=WS_CLOSE_TIMEOUT_S) except asyncio.TimeoutError: # No close-frame echo in time; websockets force-closes the transport on cancel. logger.debug( "[%s] WS close handshake exceeded %.1fs — dropping connection", self._adapter.name, WS_CLOSE_TIMEOUT_S, ) except Exception: pass class MediaSendHandler(ABC): """Media send strategy: subclasses provide acquire_file() and build_msg_body(); handle() runs the shared flow (check ws → cancel notifier → validate → COS upload → lock → dispatch).""" @abstractmethod async def acquire_file(self, adapter: "YuanbaoAdapter", **kwargs: Any) -> Tuple[bytes, str, str]: """Return (file_bytes, filename, content_type); raise ValueError when unobtainable.""" @abstractmethod def build_msg_body(self, upload_result: dict, **kwargs: Any) -> list: """Build platform-specific MsgBody list from COS upload result.""" def needs_cos_upload(self) -> bool: """Override to return False for non-COS media (e.g. sticker).""" return True async def handle( self, adapter: "YuanbaoAdapter", chat_id: str, reply_to: Optional[str] = None, caption: Optional[str] = None, **kwargs: Any, ) -> "SendResult": """Template method: shared media send flow.""" conn = adapter._connection sender = adapter._outbound.sender if conn.ws is None: return SendResult(success=False, error="Not connected", retryable=True) adapter._outbound.cancel_slow_notifier(chat_id) try: file_bytes, filename, content_type = await self.acquire_file(adapter, **kwargs) if self.needs_cos_upload(): # Stickers (TIMFaceElem) carry no bytes — validating them would yield "Empty file". validation_err = MessageSender.validate_media(file_bytes, filename, adapter.MEDIA_MAX_SIZE_MB) if validation_err: return SendResult(success=False, error=validation_err) 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, ) 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, ) else: msg_body = self.build_msg_body({}, **kwargs) if caption: msg_body.append({"msg_type": "TIMTextElem", "msg_content": {"text": caption}}) gc = kwargs.get("group_code", "") return await sender.dispatch_msg_body(chat_id, msg_body, reply_to, group_code=gc) except ValueError as ve: return SendResult(success=False, error=str(ve)) except Exception as exc: logger.error("[%s] %s.handle() failed: %s", adapter.name, type(self).__name__, exc, exc_info=True) return SendResult(success=False, error=str(exc) or type(exc).__name__) class _ImageHandler(MediaSendHandler): """Shared TIMImageElem body builder for image handlers.""" def build_msg_body(self, upload_result, **kwargs): return build_image_msg_body( url=upload_result["url"], uuid=kwargs["file_uuid"], filename=kwargs["filename"], size=upload_result["size"], width=upload_result.get("width", 0), height=upload_result.get("height", 0), mime_type=kwargs["content_type"], ) class ImageUrlHandler(_ImageHandler): """Strategy: send image from a URL (download → COS → TIMImageElem).""" async def acquire_file(self, adapter, **kwargs): image_url: str = kwargs["image_url"] logger.info("[%s] ImageUrlHandler: downloading %s", adapter.name, image_url) file_bytes, content_type = await media_download_url( image_url, max_size_mb=adapter.MEDIA_MAX_SIZE_MB, ) if not content_type or content_type == "application/octet-stream": path_part = image_url.split("?")[0] content_type = guess_mime_type(path_part) or "image/jpeg" filename = os.path.basename(image_url.split("?")[0]) or "image.jpg" return file_bytes, filename, content_type class ImageFileHandler(_ImageHandler): """Strategy: send image from a local file path (read → COS → TIMImageElem).""" async def acquire_file(self, adapter, **kwargs): image_path: str = kwargs["image_path"] if not os.path.isfile(image_path): raise ValueError(f"File not found: {image_path}") logger.info("[%s] ImageFileHandler: reading %s", adapter.name, image_path) with open(image_path, "rb") as f: file_bytes = f.read() filename = os.path.basename(image_path) or "image.jpg" content_type = guess_mime_type(filename) or "image/jpeg" return file_bytes, filename, content_type class DocumentHandler(MediaSendHandler): """Strategy: send local file/document (read → COS → TIMFileElem).""" async def acquire_file(self, adapter, **kwargs): file_path: str = kwargs["file_path"] if not os.path.isfile(file_path): raise ValueError(f"File not found: {file_path}") logger.info("[%s] DocumentHandler: reading %s", adapter.name, file_path) with open(file_path, "rb") as f: file_bytes = f.read() filename = kwargs.get("filename") or os.path.basename(file_path) or "document" content_type = guess_mime_type(filename) or "application/octet-stream" return file_bytes, filename, content_type def build_msg_body(self, upload_result, **kwargs): return build_file_msg_body( url=upload_result["url"], filename=kwargs["filename"], uuid=kwargs["file_uuid"], size=upload_result["size"], ) class StickerHandler(MediaSendHandler): """Strategy: send sticker/emoji (TIMFaceElem, no COS upload needed).""" def needs_cos_upload(self) -> bool: return False async def acquire_file(self, adapter, **kwargs): return b"", "sticker", "application/octet-stream" # no file bytes needed def build_msg_body(self, upload_result, **kwargs): from gateway.platforms.yuanbao_sticker import ( get_sticker_by_name, get_random_sticker, build_face_msg_body, build_sticker_msg_body, ) sticker_name = kwargs.get("sticker_name") face_index = kwargs.get("face_index") if sticker_name is not None: sticker = get_sticker_by_name(sticker_name) if sticker is None: raise ValueError(f"Sticker not found: {sticker_name!r}") return build_sticker_msg_body(sticker) elif face_index is not None: return build_face_msg_body(face_index=face_index) else: sticker = get_random_sticker() return build_sticker_msg_body(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"") if biz_data and isinstance(biz_data, bytes): return decode_rsp(biz_data) return empty except asyncio.TimeoutError: logger.warning("[%s] %s timeout: group=%s", adapter.name, label, group_code) return None 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]: """Query group info via WS (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]: """Query group member list via WS; 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 def _cancel_all(tasks: Dict[str, asyncio.Task]) -> None: """Cancel every unfinished task in *tasks* and clear the dict.""" for task in list(tasks.values()): if not task.done(): task.cancel() tasks.clear() class HeartbeatManager: """Reply heartbeat lifecycle: RUNNING every 2s, auto-FINISH after 30s idle, explicit stop.""" def __init__(self, adapter: "YuanbaoAdapter") -> None: self._adapter = adapter self._reply_heartbeat_tasks: Dict[str, asyncio.Task] = {} self._reply_hb_last_active: Dict[str, float] = {} 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: 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) status_name = "RUNNING" if heartbeat_val == WS_HEARTBEAT_RUNNING else "FINISH" logger.debug( "[%s] Reply heartbeat %s sent: chat=%s", adapter.name, status_name, chat_id, ) except Exception as exc: logger.debug("[%s] send_heartbeat_once failed: %s", adapter.name, exc) async def start(self, chat_id: str) -> None: """Start or renew the Reply Heartbeat periodic sender (RUNNING, every 2s).""" adapter = self._adapter conn = adapter._connection if conn.ws is None or not adapter._bot_id: return self._reply_hb_last_active[chat_id] = time.time() existing = self._reply_heartbeat_tasks.get(chat_id) if existing and not existing.done(): return self._reply_heartbeat_tasks[chat_id] = asyncio.create_task( self._worker(chat_id), name=f"yuanbao-reply-hb-{chat_id}", ) async def _worker(self, chat_id: str) -> None: """Send RUNNING every 2s; after 30s without renewal (or WS loss) send FINISH and exit.""" try: await self.send_heartbeat_once(chat_id, WS_HEARTBEAT_RUNNING) while True: await asyncio.sleep(REPLY_HEARTBEAT_INTERVAL_S) last_active = self._reply_hb_last_active.get(chat_id, 0) if time.time() - last_active > REPLY_HEARTBEAT_TIMEOUT_S: break conn = self._adapter._connection if conn.ws is None: break await self.send_heartbeat_once(chat_id, WS_HEARTBEAT_RUNNING) except asyncio.CancelledError: cancelled = True except Exception: cancelled = False else: cancelled = False finally: if not cancelled: await self.send_heartbeat_once(chat_id, WS_HEARTBEAT_FINISH) self._reply_heartbeat_tasks.pop(chat_id, None) self._reply_hb_last_active.pop(chat_id, None) async def stop(self, chat_id: str, send_finish: bool = True) -> None: """Stop Reply Heartbeat 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 if send_finish: await self.send_heartbeat_once(chat_id, WS_HEARTBEAT_FINISH) async def close(self) -> None: """Cancel all reply heartbeat tasks.""" _cancel_all(self._reply_heartbeat_tasks) self._reply_hb_last_active.clear() class SlowResponseNotifier: """Per-chat timer that sends a courtesy 'please wait' after SLOW_RESPONSE_TIMEOUT_S without a reply.""" def __init__(self, adapter: "YuanbaoAdapter", sender: "MessageSender") -> None: self._adapter = adapter self._sender = sender self._tasks: Dict[str, asyncio.Task] = {} async def start(self, chat_id: str) -> None: """Start a delayed task that notifies the user when the agent is slow.""" self.cancel(chat_id) self._tasks[chat_id] = asyncio.create_task(self._notifier(chat_id), name=f"yuanbao-slow-resp-{chat_id}") async def _notifier(self, chat_id: str) -> None: """Wait SLOW_RESPONSE_TIMEOUT_S, then push a 'please wait' message.""" try: await asyncio.sleep(SLOW_RESPONSE_TIMEOUT_S) logger.info( "[%s] Agent response exceeded %ds for %s, sending wait notice", self._adapter.name, int(SLOW_RESPONSE_TIMEOUT_S), chat_id, ) await self._sender.send_text_chunk(chat_id, SLOW_RESPONSE_MESSAGE) except asyncio.CancelledError: pass except Exception as exc: logger.debug("[%s] Slow-response notifier failed: %s", self._adapter.name, exc) def cancel(self, chat_id: str) -> None: """Cancel the pending slow-response notifier for *chat_id*, if any.""" task = self._tasks.pop(chat_id, None) if task and not task.done(): task.cancel() async def close(self) -> None: """Cancel all slow-response tasks.""" _cancel_all(self._tasks) class MessageSender: """Outbound dispatcher: per-chat locks (serial ordering), chunked text with retry, C2C/group encoding, media handler strategies, and send_direct for the send_message tool.""" IMAGE_EXTS: ClassVar[frozenset] = frozenset({".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}) CHAT_DICT_MAX_SIZE: ClassVar[int] = 1000 # Max distinct chat IDs in _chat_locks 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(), } def get_chat_lock(self, chat_id: str) -> asyncio.Lock: """Return (or create) a per-chat-id lock with safe LRU eviction.""" if chat_id in self._chat_locks: self._chat_locks.move_to_end(chat_id) return self._chat_locks[chat_id] if len(self._chat_locks) >= self.CHAT_DICT_MAX_SIZE: evicted = False for key in list(self._chat_locks): if not self._chat_locks[key].locked(): self._chat_locks.pop(key) evicted = True break if not evicted: self._chat_locks.pop(next(iter(self._chat_locks))) self._chat_locks[chat_id] = asyncio.Lock() return self._chat_locks[chat_id] async def send_text( self, chat_id: str, content: str, reply_to: Optional[str] = None, group_code: str = "", ) -> "SendResult": """Send text message with auto-chunking and per-chat-id ordering guarantee.""" adapter = self._adapter conn = adapter._connection if conn.ws is None: return SendResult(success=False, error="Not connected", retryable=True) if self._on_send_start: self._on_send_start(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) logger.info( "[%s] truncate_message: input=%d chars, max=%d, output=%d chunk(s) sizes=%s", adapter.name, len(content_to_send), adapter.MAX_TEXT_CHUNK, len(chunks), [len(c) for c in chunks], ) for i, chunk in enumerate(chunks): r_to = reply_to if i == 0 else None result = await self.send_text_chunk(chat_id, chunk, r_to, 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 return SendResult(success=True) async def send_media( self, chat_id: str, handler_name: str, reply_to: Optional[str] = None, caption: Optional[str] = None, **kwargs: Any, ) -> "SendResult": """Dispatch media send to the named handler strategy.""" handler = self._media_handlers.get(handler_name) if handler is None: return SendResult(success=False, error=f"Unknown media handler: {handler_name!r}") return await handler.handle(self._adapter, chat_id, reply_to=reply_to, caption=caption, **kwargs) async def send_direct( self, chat_id: str, message: str, media_files: Optional[List[Tuple[str, bool]]] = None, ) -> Dict[str, Any]: """send_message tool entry: text first, then each media file by extension, on the running adapter.""" adapter = self._adapter last_result: Optional["SendResult"] = None if message.strip(): last_result = await adapter.send(chat_id, message) if not last_result.success: return {"error": f"Yuanbao send failed: {last_result.error}"} for media_path, _is_voice in media_files or []: ext = Path(media_path).suffix.lower() if ext in self.IMAGE_EXTS: last_result = await adapter.send_image_file(chat_id, media_path) else: last_result = await adapter.send_document(chat_id, media_path) if not last_result.success: return {"error": f"Yuanbao media send failed: {last_result.error}"} if last_result is None: return {"error": "No deliverable text or media remained after processing"} return { "success": True, "platform": "yuanbao", "chat_id": chat_id, "message_id": last_result.message_id if last_result else None, } async def dispatch_msg_body( self, chat_id: str, msg_body: list, reply_to: Optional[str] = None, group_code: str = "", ) -> "SendResult": """Lock + dispatch an arbitrary MsgBody to C2C or group.""" async with self.get_chat_lock(chat_id): result = await self._send_msg_body(chat_id, msg_body, reply_to, group_code) if result.get("success"): return SendResult(success=True, message_id=result.get("msg_key")) return SendResult(success=False, error=result.get("error", "Unknown error")) async def send_text_chunk( self, chat_id: str, text: str, reply_to: Optional[str] = None, retry: int = 3, group_code: str = "", ) -> "SendResult": """Send a single text chunk with retry (exponential backoff: 1s, 2s, 4s).""" adapter = self._adapter last_error: str = "Unknown error" for attempt in range(retry): try: if chat_id.startswith("group:"): msg_body = self._build_msg_body_with_mentions(text, chat_id[len("group:"):]) else: msg_body = [{"msg_type": "TIMTextElem", "msg_content": {"text": text}}] raw = await self._send_msg_body(chat_id, msg_body, reply_to, group_code) if raw.get("success"): return SendResult(success=True, message_id=raw.get("msg_key")) last_error = raw.get("error", "Unknown error") logger.warning( "[%s] send_text_chunk attempt %d/%d failed: %s", adapter.name, attempt + 1, retry, last_error, ) except Exception as exc: last_error = str(exc) logger.warning( "[%s] send_text_chunk attempt %d/%d exception: %s", adapter.name, attempt + 1, retry, last_error, ) if attempt < retry - 1: await asyncio.sleep(2 ** attempt) logger.error( "[%s] send_text_chunk max retries (%d) exceeded. Last error: %s", adapter.name, retry, last_error, ) return SendResult(success=False, error=f"Max retries exceeded: {last_error}") async def _send_msg_body(self, chat_id: str, msg_body: list, reply_to: Optional[str], group_code: str) -> dict: """Route a MsgBody to group (``group:``) or C2C (``direct:`` / bare account).""" if chat_id.startswith("group:"): return await self.send_group_msg_body(chat_id[len("group:"):], msg_body, reply_to) return await self.send_c2c_msg_body(chat_id.removeprefix("direct:"), msg_body, group_code=group_code) # @nickname bounded by whitespace / line edges _AT_USER_RE = re.compile(r'(?:(?<=\s)|(?<=^))@(\S+?)(?=\s|$)', re.MULTILINE) def _build_msg_body_with_mentions(self, text: str, group_code: str) -> list: """Parse @nickname patterns and build mixed TIMTextElem + TIMCustomElem msg_body.""" members: list = [] cached = self._adapter._member_cache.get(group_code) if cached: ts, member_list = cached if time.time() - ts < self._adapter.MEMBER_CACHE_TTL_S: members = member_list else: del self._adapter._member_cache[group_code] if not members: return [{"msg_type": "TIMTextElem", "msg_content": {"text": text}}] nickname_to_uid = {} for m in members: nick = m.get("nickname") or m.get("nick_name") or "" uid = m.get("user_id") or "" if nick and uid: nickname_to_uid[nick.lower()] = (nick, uid) msg_body: list = [] last_idx = 0 for match in self._AT_USER_RE.finditer(text): start = match.start() if start > last_idx: seg = text[last_idx:start].strip() if seg: msg_body.append({"msg_type": "TIMTextElem", "msg_content": {"text": seg}}) nickname = match.group(1) entry = nickname_to_uid.get(nickname.lower()) if entry: real_nick, uid = entry msg_body.append({ "msg_type": "TIMCustomElem", "msg_content": { "data": json.dumps({"elem_type": 1002, "text": f"@{real_nick}", "user_id": uid}), }, }) else: msg_body.append({"msg_type": "TIMTextElem", "msg_content": {"text": f"@{nickname}"}}) last_idx = match.end() if last_idx < len(text): tail = text[last_idx:].strip() if tail: msg_body.append({"msg_type": "TIMTextElem", "msg_content": {"text": tail}}) if not msg_body: msg_body.append({"msg_type": "TIMTextElem", "msg_content": {"text": text}}) return msg_body async def send_c2c_msg_body(self, to_account: str, msg_body: list, group_code: str = "") -> dict: """Send C2C message with arbitrary MsgBody.""" req_id = f"c2c_{next_seq_no()}" encoded = encode_send_c2c_message( to_account=to_account, msg_body=msg_body, from_account=self._adapter._bot_id or "", msg_id=req_id, group_code=group_code, ) return await self._dispatch_encoded(self._adapter, encoded, req_id) async def send_group_msg_body(self, group_code: str, msg_body: list, reply_to: Optional[str] = None) -> dict: """Send group message with arbitrary MsgBody.""" req_id = f"grp_{next_seq_no()}" encoded = encode_send_group_message( group_code=group_code, msg_body=msg_body, from_account=self._adapter._bot_id or "", msg_id=req_id, ref_msg_id=reply_to or "", ) return await self._dispatch_encoded(self._adapter, encoded, req_id) @staticmethod async def _dispatch_encoded(adapter: "YuanbaoAdapter", encoded: bytes, req_id: str) -> dict: """Send pre-encoded bytes via WS and return a normalised result dict.""" try: response = await adapter._connection.send_biz_request(encoded, req_id=req_id) return {"success": True, "msg_key": response.get("msg_id", "")} except asyncio.TimeoutError: return {"success": False, "error": f"Request timeout after {DEFAULT_SEND_TIMEOUT}s"} except Exception as exc: return {"success": False, "error": str(exc)} @staticmethod def validate_media(file_bytes: Optional[bytes], filename: str, max_size_mb: int = 20) -> Optional[str]: """Pre-upload check; returns an error description or None.""" if file_bytes is None or len(file_bytes) == 0: return f"Empty file: {filename}" max_bytes = max_size_mb * 1024 * 1024 if len(file_bytes) > max_bytes: size_mb = len(file_bytes) / 1024 / 1024 return f"File too large: {filename} ({size_mb:.1f}MB > {max_size_mb}MB)" return None @staticmethod def truncate_message( content: str, max_length: int = 4000, len_fn: Optional[Callable[[str], int]] = None, ) -> List[str]: """Table/fence-aware chunking via MarkdownProcessor, stripping ``(1/3)`` page indicators.""" _len = len_fn or len if _len(content) <= max_length: return [content] chunks = MarkdownProcessor.chunk_markdown_text(content, max_length, len_fn=len_fn) chunks = [_INDICATOR_RE.sub('', c) for c in chunks] return chunks if chunks else [content] @staticmethod def strip_cron_wrapper(content: str) -> str: """Strip scheduler cron header/footer wrapper for cleaner Yuanbao output.""" if not content.startswith("Cronjob Response: "): return content divider = "\n-------------\n\n" footer_prefix = '\n\nTo stop or manage this job, send me a new message (e.g. "stop reminder ' divider_pos = content.find(divider) footer_pos = content.rfind(footer_prefix) if divider_pos < 0 or footer_pos < 0 or footer_pos <= divider_pos: return content header = content[:divider_pos] if "\n(job_id: " not in header: return content body_start = divider_pos + len(divider) body = content[body_start:footer_pos].strip() return body or content async def close(self) -> None: self._chat_locks.clear() class OutboundManager: """Composes MessageSender, HeartbeatManager and SlowResponseNotifier and wires their coordination hooks.""" CHAT_DICT_MAX_SIZE: ClassVar[int] = MessageSender.CHAT_DICT_MAX_SIZE def __init__(self, adapter: "YuanbaoAdapter") -> None: self._adapter = adapter self.sender: MessageSender = MessageSender(adapter) self.heartbeat: HeartbeatManager = HeartbeatManager(adapter) self.slow_notifier: SlowResponseNotifier = SlowResponseNotifier(adapter, self.sender) self.sender._on_send_start = self._handle_send_start self.sender._on_send_finish = self._handle_send_finish def _handle_send_start(self, chat_id: str) -> None: """Before a send: cancel the slow-response notifier.""" self.slow_notifier.cancel(chat_id) 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) def cancel_slow_notifier(self, chat_id: str) -> None: self.slow_notifier.cancel(chat_id) def get_chat_lock(self, chat_id: str) -> asyncio.Lock: return self.sender.get_chat_lock(chat_id) @property def _chat_locks(self) -> collections.OrderedDict: return self.sender._chat_locks async def close(self) -> None: """Shut down all sub-managers.""" await self.sender.close() await self.heartbeat.close() await self.slow_notifier.close() class YuanbaoAdapter(BasePlatformAdapter): """Yuanbao AI Bot adapter backed by a persistent WebSocket connection.""" PLATFORM = Platform.YUANBAO MAX_TEXT_CHUNK: int = 4000 # Yuanbao single message character limit splits_long_messages = True # send() auto-chunks via truncate_message(MAX_TEXT_CHUNK) MEDIA_MAX_SIZE_MB: int = 50 # Max media file size in MB for upload validation _active_instance: ClassVar[Optional["YuanbaoAdapter"]] = None @classmethod def get_active(cls) -> Optional["YuanbaoAdapter"]: """Return the currently connected YuanbaoAdapter, or None.""" return cls._active_instance @classmethod def set_active(cls, adapter: Optional["YuanbaoAdapter"]) -> None: """Register (or clear) the active adapter instance.""" cls._active_instance = adapter def __init__(self, config: PlatformConfig, **kwargs: Any) -> None: super().__init__(config, Platform.YUANBAO) _extra = config.extra or {} self._app_key: str = (_extra.get("app_id") or "").strip() self._app_secret: str = (_extra.get("app_secret") or "").strip() self._bot_id: Optional[str] = _extra.get("bot_id") or None self._ws_url: str = (_extra.get("ws_url") or DEFAULT_WS_GATEWAY_URL).strip() self._api_domain: str = (_extra.get("api_domain") or DEFAULT_API_DOMAIN).rstrip("/") self._route_env: str = (_extra.get("route_env") or "").strip() # Media resolve concurrency clamped to [min, max] so a bad config can't hammer the backend. try: _raw_concurrency = int(_extra.get("media_resolve_concurrency", _DEFAULT_RESOLVE_CONCURRENCY)) except (TypeError, ValueError): _raw_concurrency = _DEFAULT_RESOLVE_CONCURRENCY self.media_resolve_concurrency: int = max(_MIN_RESOLVE_CONCURRENCY, min(_MAX_RESOLVE_CONCURRENCY, _raw_concurrency)) self._connection: ConnectionManager = ConnectionManager(self) self._outbound: OutboundManager = OutboundManager(self) self._inbound_tasks: set[asyncio.Task] = set() # cancelled by disconnect() self._background_tasks: set[asyncio.Task] = set() # keeps fire-and-forget tasks alive # group_code -> (updated_ts, members); used by @mention resolution, stale after MEMBER_CACHE_TTL_S self._member_cache: Dict[str, Tuple[float, list]] = {} self.MEMBER_CACHE_TTL_S: float = 300.0 self._dedup = MessageDeduplicator(ttl_seconds=300) # WS reconnect / network jitter self._group_queues: Dict[str, asyncio.Queue] = {} # session_key → sequential dispatch queue # Recall support: msg_id/text being processed per session_key (RecallGuardMiddleware), plus a # bounded msg_id → content cache for content-match redaction when rows lack a message_id. self._processing_msg_ids: Dict[str, str] = {} self._processing_msg_texts: Dict[str, str] = {} self._msg_content_cache: Dict[str, str] = {} def _policy(kind: str) -> tuple[str, list[str]]: policy = (_extra.get(f"{kind}_policy") or _yb_secret(f"YUANBAO_{kind.upper()}_POLICY") or "pairing").strip().lower() raw = _extra.get(f"{kind}_allow_from") or _yb_secret(f"YUANBAO_{kind.upper()}_ALLOW_FROM", "") return policy, [x.strip() for x in raw.split(",") if x.strip()] dm_policy, dm_allow_from = _policy("dm") 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 "" ) self._auto_sethome_done: bool = bool(_existing_home) and not _existing_home.startswith("group:") def _track_task(self, task: asyncio.Task) -> asyncio.Task: """Register a fire-and-forget task so it won't be GC'd prematurely.""" self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) return task @property def enforces_own_access_policy(self) -> bool: """Yuanbao gates DM/group access at intake via dm_policy/group_policy.""" return True def _sender_may_designate_home(self, ctx: InboundContext) -> bool: """Sender may persist YUANBAO_HOME_CHANNEL: strict allowlist, open opt-in, or pairing-approved (intake-only pairing forwards are excluded).""" policy: AccessPolicy = self._access_policy sender = str(ctx.from_account or "").strip() if not sender: return False if ctx.chat_type == "dm": if policy.is_dm_allowed(sender): return True if policy.dm_policy == "pairing": 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 return False async def connect(self, *, is_reconnect: bool = False) -> bool: ok = await self._connection.open() if ok: self._wire_plugin_handlers(None) # plugin-registered native handlers return ok async def disconnect(self) -> None: """Cancel background tasks and close the WebSocket connection.""" if YuanbaoAdapter._active_instance is self: YuanbaoAdapter.set_active(None) self._running = False self._mark_disconnected() self._release_platform_lock() await self._connection.close() await self._outbound.close() for task in list(self._inbound_tasks): if not task.done(): task.cancel() self._inbound_tasks.clear() self._group_queues.clear() logger.info("[%s] Disconnected", self.name) async def send( self, chat_id: str, content: str, reply_to: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, group_code: str = "", ) -> SendResult: return await self._outbound.sender.send_text(chat_id, content, reply_to, group_code=group_code) async def get_chat_info(self, chat_id: str) -> Dict[str, Any]: """Chat metadata from the chat_id prefix (``group:`` → group, else dm).""" if chat_id.startswith("group:"): return {"name": chat_id, "type": "group"} return {"name": chat_id, "type": "dm"} async def send_typing(self, chat_id: str, metadata: Optional[dict] = None) -> None: """Start the RUNNING heartbeat.""" try: await self._outbound.heartbeat.start(chat_id) except Exception: pass async def stop_typing(self, chat_id: str) -> None: """Stop RUNNING without FINISH — send() emits FINISH after delivery so ordering is RUNNING… → message → FINISH.""" try: await self._outbound.heartbeat.stop(chat_id, send_finish=False) except Exception: pass async def _process_message_background(self, event, session_key: str) -> None: """Wrap base class processing with a slow-response notifier.""" chat_id = event.source.chat_id await self._outbound.start_slow_notifier(chat_id) try: await super()._process_message_background(event, session_key) finally: self._outbound.cancel_slow_notifier(chat_id) # Clear RecallGuard tracking only if our msg_id is still current: a concurrent message # may have overwritten it (the drain task then owns it), and id-less events never # wrote one, so they must never pop another message's entry. msg_id = event.message_id if msg_id and self._processing_msg_ids.get(session_key) == msg_id: self._processing_msg_ids.pop(session_key, None) self._processing_msg_texts.pop(session_key, None) async def query_group_info(self, group_code: str) -> Optional[dict]: return await self._group_query.query_group_info_raw(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) DM_MAX_CHARS = 10000 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.""" if not self._access_policy.is_dm_allowed(user_id): return SendResult(success=False, error="DM access denied for this user") if len(text) > self.DM_MAX_CHARS: text = text[:self.DM_MAX_CHARS] + "\n...(truncated)" chat_id = f"direct:{user_id}" return await self.send(chat_id, text, group_code=group_code) # Media sends delegate to MessageSender.send_media via the named handler strategy. async def send_image( self, chat_id: str, image_url: str, caption: Optional[str] = None, reply_to: Optional[str] = None, metadata: Optional[dict] = None, **kwargs: Any, ) -> SendResult: return await self._outbound.sender.send_media( chat_id, "image_url", reply_to=reply_to, caption=caption, image_url=image_url, **kwargs, ) async def send_image_file( self, chat_id: str, image_path: str, caption: Optional[str] = None, reply_to: Optional[str] = None, metadata: Optional[dict] = None, **kwargs: Any, ) -> SendResult: return await self._outbound.sender.send_media( chat_id, "image_file", reply_to=reply_to, caption=caption, image_path=image_path, **kwargs, ) async def send_sticker( self, chat_id: str, sticker_name: Optional[str] = None, face_index: Optional[int] = None, reply_to: Optional[str] = None, **kwargs: Any, ) -> SendResult: return await self._outbound.sender.send_media( chat_id, "sticker", reply_to=reply_to, sticker_name=sticker_name, face_index=face_index, **kwargs, ) async def send_document( self, chat_id: str, file_path: str, filename: Optional[str] = None, caption: Optional[str] = None, reply_to: Optional[str] = None, metadata: Optional[dict] = None, **kwargs: Any, ) -> SendResult: return await self._outbound.sender.send_media( chat_id, "document", reply_to=reply_to, caption=caption, file_path=file_path, filename=filename, **kwargs, ) async def _get_cached_token(self) -> dict: """Get the current valid sign token (using module-level cache).""" return await SignManager.get_token( self._app_key, self._app_secret, self._api_domain, route_env=self._route_env, ) # Module-level delegates kept for external importers (tools/send_message_tool, tools/yuanbao_tools). def get_active_adapter() -> Optional["YuanbaoAdapter"]: return YuanbaoAdapter.get_active() async def send_yuanbao_direct( adapter: "YuanbaoAdapter", chat_id: str, message: str, media_files: Optional[List[Tuple[str, bool]]] = None, ) -> Dict[str, Any]: return await adapter._outbound.sender.send_direct(chat_id, message, media_files)