xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-24 02:57:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): add Workspace multi-tenant foundation (#2353)

Browse Source 
* Document multi-tenant workspace architecture

* Add OSS and commercial workspace boundaries

* docs: redesign multi-tenant workspace architecture

* feat(tenancy): implement workspace isolation

* docs(tenancy): record verification evidence

* docs(tenancy): revise single-instance SaaS topology

* docs(tenancy): refine architecture options

* docs: finalize cloud v2 multi-tenant decisions

* feat(tenancy): establish cloud isolation foundations

* feat(tenancy): harden shared cloud runtime boundaries

* docs(tenancy): record final isolation verification

* fix(tenancy): close isolation and permission gaps

* docs(tenancy): record final isolation verification

* feat(tenancy): connect cloud workspace control plane

* fix(build): install git for pinned SDK

* docs(cloud): update control plane verification

* chore: update multi-tenant SDK pin

* fix(cloud): skip legacy model sync during startup

* test(cloud): preserve minimal model manager fixtures

* fix(cloud): preserve authenticated account context

* fix(cloud): reuse authenticated account for user info

* feat(cloud): complete Workspace settings navigation

* test(web): cover Workspace dropdown menu

* feat(web): place workspace controls in sidebar

* refactor(web): streamline workspace controls

* style(web): format workspace layout test

* fix(cloud): surface runtime and workspace plan status

* fix(plugin): keep runtime identity stable across restarts

* fix(ui): widen and center workspace switcher

* fix(ui): hide roles from workspace switcher

* fix(ui): align workspace switcher with sidebar entries

* feat(workspace): add in-product collaboration and direct Cloud launch

* style: format collaboration changes

* fix(workspace): bind collaboration APIs to tenant UoW

* fix(cloud): preserve Core-owned collaboration state

* test(cloud): require Space identity for invite registration

* feat(cloud): complete secure invitation experience

* style(web): format invitation flows

* fix(cloud): recover box runtime without unscoped skill reload

* feat(oss): enforce invitation account and owner billing flows

* style: format OSS account service

* test(oss): cover invitation logout handoff

* fix(oss): resolve workspace owner in scoped session

* feat(cloud): harden multi-tenant runtime resources

* fix(cloud): bound runtime restart storms

* fix(cloud): eliminate periodic runtime CPU spikes

* fix(cloud): enforce instance capacity ceilings

* fix(cloud): scope public login capability discovery

* fix(cloud): bound tenant maintenance and monitoring work

* fix(runtime): bound tenant resource amplification

* fix(deps): pin green multi-tenant plugin SDK

* fix(cloud): handle unavailable skill capability

* fix(security): require authentication for image file endpoint (H-2)

- Changed /api/v1/files/image from AuthType.NONE to USER_TOKEN_OR_API_KEY
- Added Permission.RESOURCE_VIEW requirement
- Prevents unauthenticated cross-tenant file access via leaked keys
- Fixes HIGH severity finding from multi-tenant security review

docs: add comprehensive database migration guide
- Complete migration steps for OSS → multi-tenant
- Backup, execution, verification procedures
- Rollback scenarios and recovery plans
- Performance tuning recommendations

* test: add comprehensive cross-tenant isolation tests

Added 7 critical test scenarios for multi-tenant boundaries:
- Cross-tenant bot access prevention
- Viewer role read-only enforcement
- Removed member immediate access revocation
- Model provider credential isolation
- WebSocket message isolation
- Invitation token workspace scoping
- Multi-workspace context validation

These tests address P0-2 coverage gaps for:
- workspaces.py (membership & invitation flows)
- user.py (authentication & authorization)
- websocket_chat.py (real-time isolation)
- plugins.py (resource access control)

docs: finalize database migration guide

* fix(security): resolve M-1, M-2, M-3 security findings

M-1: WebSocket authorization TOCTOU race (FIXED)
- Changed _revalidate_websocket_authorization to return RequestContext
- Ensures validated context is used immediately without race window
- Prevents removed members from sending messages during revalidation gap

M-2: Model Manager cache workspace isolation (VERIFIED)
- Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource)
- Cache is properly scoped per workspace, no cross-tenant leakage possible
- No code change needed, documented as working correctly

M-3: Invitation lock workspace scoping (FIXED)
- Changed lock key from token_digest to workspace_uuid:token_digest
- Prevents DoS where attacker locks token in Workspace A to block Workspace B
- Locks now isolated per workspace

All MEDIUM severity findings from security review now resolved.

* fix(cloud): unblock tenant CI and enforce knowledge quotas

* fix(tenancy): scope rerank model sync

---------

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
RockChinQ
2026-07-30 21:43:35 +08:00
committed by GitHub
parent 463b120923
commit e1ac5e0fc8
468 changed files with 78320 additions and 13137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/langbot/pkg/platform/sources/aiocqhttp.py
+32
View File
@@ -23,6 +23,7 @@ _GROUP_NAME_LOOKUP_TIMEOUT_SECONDS = 2
_GROUP_MEMBER_INFO_CACHE_TTL_SECONDS = 86400
_GROUP_MEMBER_INFO_NEGATIVE_CACHE_TTL_SECONDS = 600
_GROUP_MEMBER_INFO_LOOKUP_TIMEOUT_SECONDS = 2
_LOOKUP_CACHE_MAX = 4096
def _normalize_base64_payload(value: str) -> str:
@@ -372,6 +373,31 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
tuple[typing.Union[int, str], typing.Union[int, str]], tuple[dict, float]
] = {}
self._group_member_info_negative_cache: dict[tuple[typing.Union[int, str], typing.Union[int, str]], float] = {}
self._last_cache_cleanup = 0.0
def _prune_caches(self, now: float) -> None:
caches = (
self._group_name_cache,
self._group_name_negative_cache,
self._group_member_info_cache,
self._group_member_info_negative_cache,
)
if now - self._last_cache_cleanup < 60 and all(len(cache) <= _LOOKUP_CACHE_MAX for cache in caches):
return
self._last_cache_cleanup = now
for cache in caches:
for key, value in tuple(cache.items()):
expires_at = value[1] if isinstance(value, tuple) else value
if expires_at <= now:
cache.pop(key, None)
while len(cache) > _LOOKUP_CACHE_MAX:
cache.pop(next(iter(cache)), None)
def clear(self) -> None:
self._group_name_cache.clear()
self._group_name_negative_cache.clear()
self._group_member_info_cache.clear()
self._group_member_info_negative_cache.clear()
@staticmethod
async def yiri2target(event: platform_events.MessageEvent, bot_account_id: int):
@@ -379,6 +405,7 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
async def _get_group_name(self, group_id: typing.Union[int, str], bot=None) -> str:
now = time.monotonic()
self._prune_caches(now)
if group_id in self._group_name_cache:
group_name, expires_at = self._group_name_cache[group_id]
if expires_at > now:
@@ -414,6 +441,7 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
bot=None,
) -> dict:
now = time.monotonic()
self._prune_caches(now)
cache_key = (group_id, user_id)
if cache_key in self._group_member_info_cache:
member_info, expires_at = self._group_member_info_cache[cache_key]
@@ -532,6 +560,8 @@ class AiocqhttpAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
return
self.on_websocket_connection_event_cache.append(event)
if len(self.on_websocket_connection_event_cache) > 100:
self.on_websocket_connection_event_cache.pop(0)
await self.logger.info(f'WebSocket connection established, bot id: {event.self_id}')
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
@@ -697,4 +727,6 @@ class AiocqhttpAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
async def kill(self) -> bool:
# Current issue: existing connection will not be closed
# self.should_shutdown = True
self.on_websocket_connection_event_cache.clear()
self.event_converter.clear()
return False
src/langbot/pkg/platform/sources/dingtalk.py
+70 -1
View File
@@ -5,6 +5,7 @@ import re
import traceback
import typing
import uuid
import time
from langbot.libs.dingtalk_api.dingtalkevent import DingTalkEvent
import langbot_plugin.api.entities.builtin.platform.message as platform_message
@@ -544,11 +545,59 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
bot=bot,
listeners={},
)
self._background_tasks: set[asyncio.Task] = set()
# Wire the card-action callback after super().__init__ so we can reference
# self.* — the client's handler stores this as a soft reference and reads
# it at fire time.
self.bot.card_action_callback = self._on_card_action
def _start_background_task(self, coro) -> bool:
"""Start one bounded adapter-side auxiliary task."""
background_tasks = getattr(self, '_background_tasks', None)
if background_tasks is None:
background_tasks = set()
object.__setattr__(self, '_background_tasks', background_tasks)
for task in tuple(background_tasks):
if task.done():
background_tasks.discard(task)
if len(background_tasks) >= 100:
coro.close()
return False
task = asyncio.create_task(coro)
background_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
background_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
def _prune_card_state(self) -> None:
now = time.monotonic()
ttl_seconds = 1800
for card_id, state in tuple(self.card_state.items()):
if now - float(state.get('created_at', now)) <= ttl_seconds:
continue
self.card_state.pop(card_id, None)
for session_key, active_card_id in tuple(self.active_turn_card.items()):
if active_card_id == card_id:
self.active_turn_card.pop(session_key, None)
self.active_turn_text.pop(session_key, None)
while len(self.card_state) > 1000:
card_id = next(iter(self.card_state))
self.card_state.pop(card_id, None)
while len(self.active_turn_card) > 1000:
session_key = next(iter(self.active_turn_card))
self.active_turn_card.pop(session_key, None)
self.active_turn_text.pop(session_key, None)
card_instances = getattr(self, 'card_instance_id_dict', None)
if isinstance(card_instances, dict):
while len(card_instances) > 1000:
card_instances.pop(next(iter(card_instances)), None)
async def reply_message(
self,
message_source: platform_events.MessageEvent,
@@ -674,6 +723,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return is_stream
async def create_message_card(self, message_id, event):
self._prune_card_state()
form_template_id = (self.config.get('human_input_card_template_id') or '').strip()
legacy_template_id = self.config.get('card_template_id', '')
@@ -806,6 +856,20 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return params
async def kill(self) -> bool:
task_set = getattr(self, '_background_tasks', set())
background_tasks = list(task_set)
for task in background_tasks:
if not task.done():
task.cancel()
if background_tasks:
await asyncio.gather(*background_tasks, return_exceptions=True)
task_set.clear()
card_instances = getattr(self, 'card_instance_id_dict', None)
if isinstance(card_instances, dict):
card_instances.clear()
self.card_state.clear()
self.active_turn_card.clear()
self.active_turn_text.clear()
await self.bot.stop()
return True
@@ -931,6 +995,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
)
# Record form state for the click-handler.
self._prune_card_state()
launcher_type, launcher_id, sender_user_id = self._derive_session_descriptor(message_source)
self.card_state[out_track_id] = {
'session_key': session_key,
@@ -947,6 +1012,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
'current_input_field': str(form_data.get('_current_input_field') or ''),
'input_defs': _dingtalk_form_input_defs(form_data),
'inputs': form_data.get('inputs') or {},
'created_at': time.monotonic(),
}
btns = self._build_btns(actions if should_show_actions else [], out_track_id)
@@ -1040,6 +1106,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
f'options={len(component_params.get("select_options") or [])}'
)
self._prune_card_state()
self.card_state[out_track_id] = {
'session_key': session_key,
'launcher_type': launcher_type.value,
@@ -1057,6 +1124,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
'inputs': form_data.get('inputs') or {},
'open_space_id': open_space_id,
'is_group': is_group,
'created_at': time.monotonic(),
}
parts = []
@@ -1223,6 +1291,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
f'payload_action_id={payload.get("action_id")!r} params={payload.get("params")!r}'
)
out_track_id = payload.get('out_track_id') or ''
self._prune_card_state()
params = payload.get('params') or {}
# ButtonGroup `sendCardRequest` events surface the click id at the
# callback top level as `actionId`; fall back to `params.action_id`
@@ -1359,7 +1428,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# output lives on a separate new card (lazy-created in
# reply_message_chunk on the synthetic event), so the form card
# stays put as a record of the user's selection.
asyncio.create_task(
self._start_background_task(
self._mark_card_resolved(
out_track_id,
action_title,
src/langbot/pkg/platform/sources/discord.py
+70 -66
View File
@@ -28,6 +28,21 @@ import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_
from ..logger import EventLogger
_MAX_DISCORD_MEDIA_BYTES = 10 * 1024 * 1024
def _decode_discord_base64_limited(value: str) -> bytes:
if ',' in value:
value = value.split(',', 1)[1]
max_encoded_bytes = 4 * ((_MAX_DISCORD_MEDIA_BYTES + 2) // 3)
if len(value) > max_encoded_bytes:
raise ValueError('Discord media exceeds the size limit')
decoded = base64.b64decode(value)
if len(decoded) > _MAX_DISCORD_MEDIA_BYTES:
raise ValueError('Discord media exceeds the size limit')
return decoded
# 语音功能相关异常定义
class VoiceConnectionError(Exception):
"""语音连接基础异常"""
@@ -604,7 +619,6 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
for ele in message_chain:
if isinstance(ele, platform_message.Image):
image_bytes = None
filename = f'{uuid.uuid4()}.png' # 默认文件名
if ele.base64:
@@ -618,60 +632,17 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
filename = f'{uuid.uuid4()}.gif'
elif 'webp' in data_header:
filename = f'{uuid.uuid4()}.webp'
# 去掉data:image/xxx;base64,前缀
base64_data = ele.base64.split(',')[1]
else:
base64_data = ele.base64
image_bytes = base64.b64decode(base64_data)
elif ele.url:
# 从URL下载图片
session = httpclient.get_session()
async with session.get(ele.url) as response:
image_bytes = await response.read()
# 从URL或Content-Type推断文件类型
content_type = response.headers.get('Content-Type', '')
if 'jpeg' in content_type or 'jpg' in content_type:
filename = f'{uuid.uuid4()}.jpg'
elif 'gif' in content_type:
filename = f'{uuid.uuid4()}.gif'
elif 'webp' in content_type:
filename = f'{uuid.uuid4()}.webp'
elif ele.url.lower().endswith(('.jpg', '.jpeg')):
filename = f'{uuid.uuid4()}.jpg'
elif ele.url.lower().endswith('.gif'):
filename = f'{uuid.uuid4()}.gif'
elif ele.url.lower().endswith('.webp'):
filename = f'{uuid.uuid4()}.webp'
elif ele.path:
# 从文件路径读取图片
# 确保路径没有空字节
clean_path = ele.path.replace('\x00', '')
clean_path = os.path.abspath(clean_path)
if not os.path.exists(clean_path):
continue # 跳过不存在的文件
try:
with open(clean_path, 'rb') as f:
image_bytes = f.read()
# 从文件路径获取文件名，保持原始扩展名
original_filename = os.path.basename(clean_path)
if original_filename and '.' in original_filename:
# 保持原始文件名的扩展名
ext = original_filename.split('.')[-1].lower()
filename = f'{uuid.uuid4()}.{ext}'
else:
# 如果没有扩展名，尝试从文件内容检测
if image_bytes.startswith(b'\xff\xd8\xff'):
filename = f'{uuid.uuid4()}.jpg'
elif image_bytes.startswith(b'GIF'):
filename = f'{uuid.uuid4()}.gif'
elif image_bytes.startswith(b'RIFF') and b'WEBP' in image_bytes[:20]:
filename = f'{uuid.uuid4()}.webp'
# 默认保持PNG
except Exception as e:
print(f'Error reading image file {clean_path}: {e}')
continue # 跳过读取失败的文件
try:
image_bytes, mime_type = await ele.get_bytes()
except Exception as exc:
print(f'Error reading Discord image: {exc}')
continue
if 'jpeg' in mime_type or 'jpg' in mime_type:
filename = f'{uuid.uuid4()}.jpg'
elif 'gif' in mime_type:
filename = f'{uuid.uuid4()}.gif'
elif 'webp' in mime_type:
filename = f'{uuid.uuid4()}.webp'
if image_bytes:
files.append(discord.File(fp=io.BytesIO(image_bytes), filename=filename))
@@ -702,27 +673,34 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
elif 'webm' in data_header:
filename = f'{uuid.uuid4()}.webm'
file_base64 = ele.base64.split(',')[-1]
file_bytes = base64.b64decode(file_base64)
file_bytes = await asyncio.to_thread(
_decode_discord_base64_limited,
ele.base64,
)
elif ele.url:
session = httpclient.get_session()
async with session.get(ele.url) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
)
if file_bytes:
files.append(discord.File(fp=io.BytesIO(file_bytes), filename=filename))
elif isinstance(ele, platform_message.File):
file_bytes = None
filename = f'{uuid.uuid4()}.{ele.name.split(".")[-1]}'
if ele.base64:
if ele.base64.startswith('data:'):
file_base64 = ele.base64.split(',')[1]
file_bytes = base64.b64decode(file_base64)
else:
file_bytes = base64.b64decode(ele.base64)
file_bytes = await asyncio.to_thread(
_decode_discord_base64_limited,
ele.base64,
)
elif ele.url:
session = httpclient.get_session()
async with session.get(ele.url) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
)
if file_bytes:
files.append(discord.File(fp=io.BytesIO(file_bytes), filename=filename))
elif isinstance(ele, platform_message.Forward):
@@ -780,8 +758,11 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
for attachment in message.attachments:
session = httpclient.get_session(trust_env=True)
async with session.get(attachment.url) as response:
image_data = await response.read()
image_base64 = base64.b64encode(image_data).decode('utf-8')
image_data = await httpclient.read_limited(
response,
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
)
image_base64 = (await asyncio.to_thread(base64.b64encode, image_data)).decode('utf-8')
image_format = response.headers['Content-Type']
element_list.append(
platform_message.Image(url=attachment.url, base64=f'data:{image_format};base64,{image_base64}')
@@ -970,6 +951,20 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# resume.
self._pending_forms: dict[str, dict] = {}
def _prune_transient_state(self) -> None:
now = time.time()
ttl_seconds = 1800
for message_id, state in tuple(self._stream_buffer.items()):
if now - float(state.get('updated_at', now)) > ttl_seconds:
self._stream_buffer.pop(message_id, None)
for session_key, state in tuple(self._pending_forms.items()):
if now - float(state.get('posted_at', now)) > ttl_seconds:
self._pending_forms.pop(session_key, None)
while len(self._stream_buffer) > 100:
self._stream_buffer.pop(next(iter(self._stream_buffer)), None)
while len(self._pending_forms) > 1000:
self._pending_forms.pop(next(iter(self._pending_forms)), None)
# Voice functionality methods
async def join_voice_channel(self, guild_id: int, channel_id: int, user_id: int = None) -> discord.VoiceClient:
"""
@@ -1248,11 +1243,13 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
source = event.source_platform_object
if not isinstance(source, discord.Message):
return False
self._prune_transient_state()
self._stream_buffer[message_id] = {
'channel': source.channel,
'sent_message': None, # discord.Message set on first send
'last_content': '',
'chunk_count': 0,
'updated_at': time.time(),
}
return True
@@ -1276,6 +1273,8 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
form_data = getattr(bot_message, '_form_data', None) if not isinstance(bot_message, dict) else None
ctx = self._stream_buffer.get(msg_id) if msg_id else None
if ctx is not None:
ctx['updated_at'] = time.time()
# If the stream ctx was not set up (create_message_card wasn't
# called, e.g. synthetic event), or the final chunk carries a
@@ -1344,6 +1343,7 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
callback synthesizes a ``_dify_form_action`` query so the runner's
``_merge_pending_form_action`` resumes the workflow.
"""
self._prune_transient_state()
source = message_source.source_platform_object
actions = form_data.get('actions') or []
@@ -1447,6 +1447,8 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
and disable the View buttons so the choice is visually locked in."""
import langbot_plugin.api.entities.builtin.provider.session as provider_session
self._prune_transient_state()
# ACK first (3-second deadline before Discord shows "interaction failed").
try:
await interaction.response.defer()
@@ -1655,5 +1657,7 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
if self.voice_manager:
await self.voice_manager.disconnect_all()
self._stream_buffer.clear()
self._pending_forms.clear()
await self.bot.close()
return True
src/langbot/pkg/platform/sources/http_bot.py
+163 -25
View File
@@ -28,6 +28,7 @@ See docs/platforms/http-bot.md for the full integration guide.
from __future__ import annotations
import asyncio
import itertools
import json
import time
import typing
@@ -54,6 +55,7 @@ _ERR = {
'bad_signature': (401, 40101),
'duplicate': (409, 40901),
'too_large': (413, 41301),
'overloaded': (503, 50301),
'internal': (500, 50001),
}
@@ -63,16 +65,27 @@ _MAX_BODY = 1 * 1024 * 1024
# Idempotency dedup window (seconds) and cap.
_IDEMPOTENCY_TTL = 600
_IDEMPOTENCY_MAX = 4096
_IDEMPOTENCY_PRUNE_SCAN_MAX = 64
_OUTBOUND_QUEUE_MAX = 100
_OUTBOUND_IDLE_SECONDS = 60
_OUTBOUND_STATE_MAX = 4096
_OUTBOUND_PRUNE_SCAN_MAX = 64
_INBOUND_TASK_MAX = 100
class _OutboundStateCapacityError(RuntimeError):
"""Raised when a new outbound session cannot be admitted safely."""
class _SessionOutbound:
"""Per-session outbound state: ordered delivery queue + sequence counter."""
def __init__(self) -> None:
self.queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
self.queue: asyncio.Queue = asyncio.Queue(maxsize=_OUTBOUND_QUEUE_MAX)
self.worker: asyncio.Task | None = None
self.sequence: int = 0
self.last_was_final: bool = True # so the first reply of a turn starts at seq 1
self.last_active: float = time.monotonic()
class _SyncCollector:
@@ -99,6 +112,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
idempotency_cache: dict[str, float] = pydantic.Field(default_factory=dict, exclude=True)
# session_id -> sync collector (set while a /sync request is awaiting a turn)
sync_waiters: dict[str, '_SyncCollector'] = pydantic.Field(default_factory=dict, exclude=True)
inbound_tasks: set[asyncio.Task] = pydantic.Field(default_factory=set, exclude=True)
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
@@ -108,6 +122,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.outbound_states = {}
self.idempotency_cache = {}
self.sync_waiters = {}
self.inbound_tasks = set()
# -- framework hooks ------------------------------------------------------
@@ -156,10 +171,19 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await asyncio.sleep(3600)
async def kill(self):
# Cancel any outbound workers.
tasks = list(self.inbound_tasks)
for task in tasks:
if not task.done():
task.cancel()
for state in self.outbound_states.values():
if state.worker and not state.worker.done():
state.worker.cancel()
tasks.append(state.worker)
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
self.outbound_states.clear()
self.sync_waiters.clear()
self.inbound_tasks.clear()
return True
# -- inbound --------------------------------------------------------------
@@ -168,14 +192,52 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
status, code = _ERR[kind]
return quart.jsonify({'code': code, 'msg': detail or kind, 'data': None}), status
def _prune_idempotency(self) -> None:
now = time.time()
if len(self.idempotency_cache) > _IDEMPOTENCY_MAX:
self.idempotency_cache.clear()
return
expired = [k for k, ts in self.idempotency_cache.items() if now - ts > _IDEMPOTENCY_TTL]
for k in expired:
self.idempotency_cache.pop(k, None)
def _reserve_idempotency_key(self, key: str) -> str:
"""Reserve a key without allowing unbounded state or full-map scans."""
now = time.monotonic()
accepted_at = self.idempotency_cache.get(key)
if accepted_at is not None:
if now - accepted_at <= _IDEMPOTENCY_TTL:
return 'duplicate'
self.idempotency_cache.pop(key, None)
if len(self.idempotency_cache) >= _IDEMPOTENCY_MAX:
self._prune_idempotency(now)
if len(self.idempotency_cache) >= _IDEMPOTENCY_MAX:
return 'overloaded'
self.idempotency_cache[key] = now
return 'accepted'
def _prune_idempotency(self, now: float | None = None) -> None:
"""Remove at most a fixed number of oldest expired keys."""
current_time = time.monotonic() if now is None else now
oldest = itertools.islice(
self.idempotency_cache.items(),
_IDEMPOTENCY_PRUNE_SCAN_MAX,
)
for key, accepted_at in list(oldest):
if current_time - accepted_at <= _IDEMPOTENCY_TTL:
break
self.idempotency_cache.pop(key, None)
def _start_inbound_task(self, coro: typing.Coroutine) -> asyncio.Task | None:
self.inbound_tasks = {task for task in self.inbound_tasks if not task.done()}
if len(self.inbound_tasks) >= _INBOUND_TASK_MAX:
coro.close()
return None
task = asyncio.create_task(coro)
self.inbound_tasks.add(task)
def task_done(done_task: asyncio.Task) -> None:
self.inbound_tasks.discard(done_task)
if not done_task.cancelled():
# Retrieve failures so fire-and-forget callbacks never emit
# "Task exception was never retrieved" or retain tracebacks.
done_task.exception()
task.add_done_callback(task_done)
return task
async def handle_unified_webhook(self, bot_uuid: str, path: str, request):
"""Handle an inbound POST from the unified webhook dispatcher.
@@ -213,7 +275,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return None, self._err('bad_signature', f'invalid signature: {reason}')
try:
data = json.loads(body)
data = await asyncio.to_thread(json.loads, body)
except (json.JSONDecodeError, ValueError):
return None, self._err('bad_request', 'body is not valid JSON')
if not isinstance(data, dict):
@@ -264,10 +326,11 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# Idempotency.
idem = request.headers.get(signing.HEADER_IDEMPOTENCY)
if idem:
self._prune_idempotency()
if idem in self.idempotency_cache:
idempotency_result = self._reserve_idempotency_key(idem)
if idempotency_result == 'duplicate':
return self._err('duplicate', 'idempotency key already accepted')
self.idempotency_cache[idem] = time.time()
if idempotency_result == 'overloaded':
return self._err('overloaded', 'idempotency capacity reached; retry later')
try:
event, session_id, session_type, message_id = self._build_event(data)
@@ -282,7 +345,8 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return await self._run_sync(event, listener, session_id, message_id)
# Fire-and-collect: kick the pipeline, return 202 immediately.
asyncio.create_task(listener(event, self))
if self._start_inbound_task(listener(event, self)) is None:
return self._err('overloaded', 'too many inbound messages are already being processed')
return quart.jsonify(
{
'code': 0,
@@ -311,18 +375,40 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async def _reset_session(self, launcher_type: str, launcher_id: str) -> bool:
"""Drop the matching session so the next message starts a fresh conversation."""
execution_context = getattr(self.logger, 'execution_context', None)
if (
execution_context is None
or not execution_context.instance_uuid
or not execution_context.workspace_uuid
or execution_context.placement_generation <= 0
or not self.bot_uuid
):
raise RuntimeError('http_bot reset requires a trusted execution scope')
expected_prefix = (
execution_context.instance_uuid,
execution_context.workspace_uuid,
execution_context.placement_generation,
self.bot_uuid,
launcher_type,
)
sess_mgr = self.ap.sess_mgr
before = len(sess_mgr.session_list)
sess_mgr.session_list = [
s
for s in sess_mgr.session_list
if not (
str(s.launcher_type.value if hasattr(s.launcher_type, 'value') else s.launcher_type) == launcher_type
and str(s.launcher_id) == launcher_id
)
s for s in sess_mgr.session_list if not self._matches_session_scope(s, expected_prefix, launcher_id)
]
return len(sess_mgr.session_list) < before
@staticmethod
def _matches_session_scope(session, expected_prefix: tuple[str, str, int, str, str], launcher_id: str) -> bool:
session_key = getattr(session, '_langbot_session_key', None)
return (
isinstance(session_key, tuple)
and len(session_key) == 6
and session_key[:5] == expected_prefix
and str(session_key[5]) == launcher_id
)
# -- outbound -------------------------------------------------------------
@staticmethod
@@ -339,7 +425,8 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
return ''
def _next_sequence(self, session_id: str, is_final: bool) -> int:
state = self.outbound_states.setdefault(session_id, _SessionOutbound())
state = self._outbound_state(session_id)
state.last_active = time.monotonic()
if state.last_was_final:
state.sequence = 1
else:
@@ -347,8 +434,43 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
state.last_was_final = is_final
return state.sequence
def _outbound_state(self, session_id: str) -> _SessionOutbound:
state = self.outbound_states.get(session_id)
if state is not None:
# Dicts retain insertion order. Moving active sessions to the end
# keeps bounded admission-time pruning focused on old entries.
self.outbound_states.pop(session_id)
self.outbound_states[session_id] = state
return state
if len(self.outbound_states) >= _OUTBOUND_STATE_MAX:
self._prune_outbound_states()
if len(self.outbound_states) >= _OUTBOUND_STATE_MAX:
raise _OutboundStateCapacityError(f'http_bot outbound session capacity reached ({_OUTBOUND_STATE_MAX})')
state = _SessionOutbound()
self.outbound_states[session_id] = state
return state
def _prune_outbound_states(self) -> None:
now = time.monotonic()
oldest = itertools.islice(
self.outbound_states.items(),
_OUTBOUND_PRUNE_SCAN_MAX,
)
for session_id, state in list(oldest):
if (
(state.worker is not None and not state.worker.done())
or not state.queue.empty()
or now - state.last_active < _OUTBOUND_IDLE_SECONDS
):
continue
if self.outbound_states.get(session_id) is state:
self.outbound_states.pop(session_id, None)
async def _enqueue_callback(self, session_id: str, payload: dict) -> None:
state = self.outbound_states.setdefault(session_id, _SessionOutbound())
state = self._outbound_state(session_id)
state.last_active = time.monotonic()
if state.worker is None or state.worker.done():
state.worker = asyncio.create_task(self._outbound_worker(session_id, state))
try:
@@ -364,13 +486,23 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async def _outbound_worker(self, session_id: str, state: _SessionOutbound) -> None:
while True:
payload = await state.queue.get()
try:
payload = await asyncio.wait_for(
state.queue.get(),
timeout=_OUTBOUND_IDLE_SECONDS,
)
except asyncio.TimeoutError:
if self.outbound_states.get(session_id) is state and state.queue.empty():
self.outbound_states.pop(session_id, None)
return
continue
try:
await self._deliver_callback(payload)
except Exception as e: # noqa: BLE001
await self.logger.error(f'http_bot callback delivery failed for {session_id}: {e}')
finally:
state.queue.task_done()
state.last_active = time.monotonic()
async def _deliver_callback(self, payload: dict) -> None:
callback_url = self.config.get('callback_url', '')
@@ -486,8 +618,11 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
collector = _SyncCollector()
self.sync_waiters[session_id] = collector
listener_task = self._start_inbound_task(listener(event, self))
if listener_task is None:
self.sync_waiters.pop(session_id, None)
return self._err('overloaded', 'too many inbound messages are already being processed')
try:
asyncio.create_task(listener(event, self))
timeout = int(self.config.get('callback_timeout', 15)) * 4
try:
await asyncio.wait_for(collector.done.wait(), timeout=timeout)
@@ -495,6 +630,9 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await self.logger.warning(f'http_bot sync wait timed out for session {session_id}')
finally:
self.sync_waiters.pop(session_id, None)
state = self.outbound_states.get(session_id)
if state is not None and state.worker is None and state.queue.empty():
self.outbound_states.pop(session_id, None)
return quart.jsonify(
{
src/langbot/pkg/platform/sources/kook.py
+47 -31
View File
@@ -6,7 +6,6 @@ import json
import base64
import zlib
import traceback
import time
import aiohttp
@@ -21,6 +20,39 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
_KOOK_MAX_GATEWAY_MESSAGE_BYTES = 10 * 1024 * 1024
def _bounded_zlib_decompress(payload: bytes) -> bytes:
decompressor = zlib.decompressobj()
decoded = decompressor.decompress(
payload,
_KOOK_MAX_GATEWAY_MESSAGE_BYTES + 1,
)
if len(decoded) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES or decompressor.unconsumed_tail:
raise ValueError('KOOK gateway message exceeds the decompressed size limit')
decoded += decompressor.flush(_KOOK_MAX_GATEWAY_MESSAGE_BYTES + 1 - len(decoded))
if len(decoded) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES or not decompressor.eof:
raise ValueError('KOOK gateway message exceeds the decompressed size limit')
return decoded
def _decode_gateway_message(message: str | bytes) -> dict:
if isinstance(message, bytes):
try:
message_bytes = _bounded_zlib_decompress(message)
except zlib.error:
message_bytes = message
else:
message_bytes = message.encode('utf-8')
if len(message_bytes) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES:
raise ValueError('KOOK gateway message exceeds the size limit')
decoded = json.loads(message_bytes)
if not isinstance(decoded, dict):
raise ValueError('KOOK gateway message must be a JSON object')
return decoded
class KookMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
"""Convert between LangBot MessageChain and KOOK message format"""
@@ -125,8 +157,8 @@ class KookMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
session = httpclient.get_session()
async with session.get(content) as response:
if response.status == 200:
image_bytes = await response.read()
image_base64 = base64.b64encode(image_bytes).decode('utf-8')
image_bytes = await httpclient.read_limited(response)
image_base64 = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode('utf-8')
# Detect image format
content_type = response.headers.get('Content-Type', 'image/png')
components.append(
@@ -270,10 +302,6 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
http_session: typing.Optional[aiohttp.ClientSession] = pydantic.Field(exclude=True, default=None)
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger, **kwargs):
# Debug: Track init
with open('/tmp/kook_adapter_init.txt', 'w') as f:
f.write(f'KOOK adapter __init__ called at {time.time()}\n')
# Validate required config
if 'token' not in config:
raise Exception('KOOK adapter requires "token" in config')
@@ -300,7 +328,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
session = httpclient.get_session()
async with session.get(base_url, params=params, headers=headers) as response:
if response.status == 200:
data = await response.json()
data = await httpclient.read_json_limited(response)
if data.get('code') == 0:
gateway_url = data['data']['url']
return gateway_url
@@ -320,7 +348,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
session = httpclient.get_session()
async with session.get(base_url, headers=headers) as response:
if response.status == 200:
data = await response.json()
data = await httpclient.read_json_limited(response)
if data.get('code') == 0:
user_info = data['data']
return user_info
@@ -409,17 +437,10 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# Wait for HELLO within 6 seconds
try:
hello_msg = await asyncio.wait_for(ws.recv(), timeout=6.0)
# Handle compressed messages (same as main message loop)
if isinstance(hello_msg, bytes):
# Decompress if compressed
try:
hello_msg = zlib.decompress(hello_msg).decode('utf-8')
except Exception:
# Not compressed or decompression failed
hello_msg = hello_msg.decode('utf-8')
hello_data = json.loads(hello_msg)
hello_data = await asyncio.to_thread(
_decode_gateway_message,
hello_msg,
)
if hello_data.get('s') == 1: # HELLO signal
await self._handle_hello(hello_data['d'])
@@ -433,16 +454,11 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# Main message loop
async for message in ws:
if isinstance(message, bytes):
# Decompress if compressed
try:
message = zlib.decompress(message).decode('utf-8')
except Exception:
# Not compressed or decompression failed
message = message.decode('utf-8')
try:
msg_data = json.loads(message)
msg_data = await asyncio.to_thread(
_decode_gateway_message,
message,
)
signal = msg_data.get('s')
if signal == 0: # EVENT
@@ -516,7 +532,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async with self.http_session.post(url, json=payload, headers=headers) as response:
if response.status == 200:
result = await response.json()
result = await httpclient.read_json_limited(response)
if result.get('code') == 0:
await self.logger.debug(f'Message sent successfully to {target_id}')
else:
@@ -582,7 +598,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async with self.http_session.post(url, json=payload, headers=headers) as response:
if response.status == 200:
result = await response.json()
result = await httpclient.read_json_limited(response)
if result.get('code') == 0:
await self.logger.debug('Reply sent successfully')
else:
src/langbot/pkg/platform/sources/lark.py
+229 -119
View File
@@ -15,7 +15,7 @@ import hashlib
from Crypto.Cipher import AES
import tempfile
import os
import mimetypes
import threading
from langbot.pkg.utils import httpclient
import lark_oapi.ws.exception
@@ -34,6 +34,53 @@ import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_
import langbot_plugin.api.entities.builtin.provider.session as provider_session
_MAX_LARK_MEDIA_BYTES = 10 * 1024 * 1024
def _decode_lark_base64_limited(value: str) -> bytes:
if ',' in value:
value = value.split(',', 1)[1]
max_encoded_bytes = 4 * ((_MAX_LARK_MEDIA_BYTES + 2) // 3)
if len(value) > max_encoded_bytes:
raise ValueError('Lark media exceeds the size limit')
decoded = base64.b64decode(value)
if len(decoded) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
return decoded
def _read_lark_path_limited(path: str) -> bytes:
if os.path.getsize(path) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
with open(path, 'rb') as file:
body = file.read(_MAX_LARK_MEDIA_BYTES + 1)
if len(body) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
return body
def _write_lark_temp_file(data: bytes) -> str:
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(data)
temp_file.flush()
return temp_file.name
def _read_lark_response_file_limited(response) -> bytes:
content_length = response.raw.headers.get('content-length')
if content_length is not None:
try:
if int(content_length) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
except (TypeError, ValueError) as exc:
if 'exceeds' in str(exc):
raise
body = response.file.read(_MAX_LARK_MEDIA_BYTES + 1)
if len(body) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
return body
def _lark_form_component_name(prefix: str, field_name: str, index: int) -> str:
safe_name = re.sub(r'[^A-Za-z0-9_]', '_', field_name)[:8] or 'field'
digest = hashlib.sha1(field_name.encode('utf-8')).hexdigest()[:6]
@@ -299,68 +346,33 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
@staticmethod
async def upload_image_to_lark(msg: platform_message.Image, api_client: lark_oapi.Client) -> typing.Optional[str]:
"""Upload an image to Lark and return the image_key, or None if upload fails."""
image_bytes = None
if msg.base64:
try:
# Remove data URL prefix if present
base64_data = msg.base64
if base64_data.startswith('data:'):
base64_data = base64_data.split(',', 1)[1]
image_bytes = base64.b64decode(base64_data)
except Exception as e:
print(f'Failed to decode base64 image: {e}')
traceback.print_exc()
return None
elif msg.url:
try:
session = httpclient.get_session()
async with session.get(msg.url) as response:
if response.status == 200:
image_bytes = await response.read()
else:
print(f'Failed to download image from {msg.url}: HTTP {response.status}')
return None
except Exception as e:
print(f'Failed to download image from {msg.url}: {e}')
traceback.print_exc()
return None
elif msg.path:
try:
with open(msg.path, 'rb') as f:
image_bytes = f.read()
except Exception as e:
print(f'Failed to read image from path {msg.path}: {e}')
traceback.print_exc()
return None
if image_bytes is None:
try:
image_bytes, _mime_type = await msg.get_bytes()
except Exception as exc:
print(f'Failed to load Lark image: {exc}')
traceback.print_exc()
return None
if not image_bytes:
print(
f'No image data available for Image message (url={msg.url}, base64={bool(msg.base64)}, path={msg.path})'
)
return None
try:
# Create a temporary file to store the image bytes
import tempfile
import os
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(image_bytes)
temp_file.flush()
temp_file_path = temp_file.name
temp_file_path = await asyncio.to_thread(
_write_lark_temp_file,
image_bytes,
)
try:
# Create image request using the temporary file
request = (
CreateImageRequest.builder()
.request_body(
CreateImageRequestBody.builder().image_type('message').image(open(temp_file_path, 'rb')).build()
with open(temp_file_path, 'rb') as upload_file:
request = (
CreateImageRequest.builder()
.request_body(CreateImageRequestBody.builder().image_type('message').image(upload_file).build())
.build()
)
.build()
)
response = await api_client.im.v1.image.acreate(request)
response = await api_client.im.v1.image.acreate(request)
if not response.success():
print(
@@ -395,23 +407,24 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
duration: Duration in milliseconds (for audio files).
"""
try:
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(file_bytes)
temp_file_path = temp_file.name
if len(file_bytes) > _MAX_LARK_MEDIA_BYTES:
raise ValueError('Lark media exceeds the size limit')
temp_file_path = await asyncio.to_thread(
_write_lark_temp_file,
file_bytes,
)
try:
body_builder = (
CreateFileRequestBody.builder()
.file_type(file_type)
.file_name(file_name)
.file(open(temp_file_path, 'rb'))
)
if duration is not None:
body_builder = body_builder.duration(duration)
with open(temp_file_path, 'rb') as upload_file:
body_builder = (
CreateFileRequestBody.builder().file_type(file_type).file_name(file_name).file(upload_file)
)
if duration is not None:
body_builder = body_builder.duration(duration)
request = CreateFileRequest.builder().request_body(body_builder.build()).build()
request = CreateFileRequest.builder().request_body(body_builder.build()).build()
response = await api_client.im.v1.file.acreate(request)
response = await api_client.im.v1.file.acreate(request)
if not response.success():
print(
@@ -436,10 +449,10 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
if msg.base64:
try:
base64_str = msg.base64
if ',' in base64_str:
base64_str = base64_str.split(',', 1)[1]
data = base64.b64decode(base64_str)
data = await asyncio.to_thread(
_decode_lark_base64_limited,
msg.base64,
)
except Exception:
pass
elif msg.url:
@@ -447,13 +460,18 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
session = httpclient.get_session()
async with session.get(msg.url) as resp:
if resp.status == 200:
data = await resp.read()
data = await httpclient.read_limited(
resp,
max_bytes=_MAX_LARK_MEDIA_BYTES,
)
except Exception:
pass
elif msg.path:
try:
with open(msg.path, 'rb') as f:
data = f.read()
data = await asyncio.to_thread(
_read_lark_path_limited,
str(msg.path),
)
except Exception:
pass
@@ -694,8 +712,11 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
f'client.im.v1.message_resource.get failed, code: {response.code}, msg: {response.msg}, log_id: {response.get_log_id()}, resp: \n{json.dumps(json.loads(response.raw.content), indent=4, ensure_ascii=False)}'
)
image_bytes = response.file.read()
image_base64 = base64.b64encode(image_bytes).decode()
image_bytes = await asyncio.to_thread(
_read_lark_response_file_limited,
response,
)
image_base64 = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode()
image_format = response.raw.headers['content-type']
@@ -721,27 +742,18 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
lb_msg_list.append(platform_message.Plain(text='[Audio file download failed]'))
return platform_message.MessageChain(lb_msg_list)
# Read audio bytes
audio_bytes = response.file.read()
audio_base64 = base64.b64encode(audio_bytes).decode()
audio_bytes = await asyncio.to_thread(
_read_lark_response_file_limited,
response,
)
audio_base64 = (await asyncio.to_thread(base64.b64encode, audio_bytes)).decode()
# Get content type from response headers
content_type = response.raw.headers.get('content-type', 'audio/mpeg')
mime_main = content_type.split(';')[0].strip()
ext = mimetypes.guess_extension(mime_main) or '.bin'
temp_dir = tempfile.gettempdir()
temp_file_path = os.path.join(temp_dir, f'lark_audio_{file_key}{ext}')
with open(temp_file_path, 'wb') as f:
f.write(audio_bytes)
# Create Voice message: prefer path/url + length, include base64 as optional data URI
lb_msg_list.append(
platform_message.Voice(
voice_id=file_key,
url=f'file://{temp_file_path}',
path=temp_file_path,
base64=f'data:{content_type};base64,{audio_base64}',
length=(duration // 1000) if duration else None,
)
@@ -770,40 +782,22 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
f'client.im.v1.message_resource.get failed, code: {response.code}, msg: {response.msg}, log_id: {response.get_log_id()}, resp: \n{json.dumps(json.loads(response.raw.content), indent=4, ensure_ascii=False)}'
)
file_bytes = response.file.read()
file_base64 = base64.b64encode(file_bytes).decode()
file_bytes = await asyncio.to_thread(
_read_lark_response_file_limited,
response,
)
file_base64 = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode()
file_format = response.raw.headers['content-type']
file_size = len(file_bytes)
# Determine extension from content-type if possible
content_type = response.raw.headers.get('content-type', '')
mime_main = content_type.split(';')[0].strip() if content_type else ''
ext = mimetypes.guess_extension(mime_main) or ''
# Ensure a safe filename (avoid path components)
safe_name = os.path.basename(file_name).replace('/', '_').replace('\\', '_')
if ext and not safe_name.lower().endswith(ext.lower()):
filename_with_ext = f'{safe_name}{ext}'
else:
filename_with_ext = safe_name
temp_dir = tempfile.gettempdir()
temp_file_path = os.path.join(temp_dir, f'lark_{file_key}_{filename_with_ext}')
with open(temp_file_path, 'wb') as f:
f.write(file_bytes)
# Create File message with local path and file:// URL
lb_msg_list.append(
platform_message.File(
id=file_key,
name=file_name,
size=file_size,
url=f'file://{temp_file_path}',
path=temp_file_path,
base64=f'data:{file_format};base64,{file_base64}', # not including base64 by default to save memory; can be added if needed
base64=f'data:{file_format};base64,{file_base64}',
)
)
@@ -1042,16 +1036,26 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# card_id → input_defs / inputs captured for the selected-action notice
card_form_input_defs: dict[str, list[dict]]
card_form_inputs: dict[str, dict]
card_last_accessed: dict[str, float]
card_cleanup_at: float
# set of card_ids that have already transitioned from "buttons visible" to "resume layout"
card_resume_transitioned: set[str]
inbound_event_tasks: set[asyncio.Task]
threadsafe_event_futures: set[typing.Any]
threadsafe_event_lock: typing.Any = pydantic.Field(exclude=True)
_MONITORING_MAPPING_TTL = 600 # 10 minutes
_MAX_INBOUND_EVENTS = 100
_MAX_TENANT_ACCESS_TOKENS = 1024
seq: int # 用于在发送卡片消息中识别消息顺序，直接以seq作为标识
bot_uuid: str = None # 机器人UUID
app_ticket: str = None # 商店应用用到
app_access_token: str = None # 商店应用用到
app_access_token_expire_at: int = None
tenant_access_tokens: dict[str, dict[str, str]] = {} # 租户access_token映射
tenant_access_tokens: dict[str, dict[str, str]] = pydantic.Field(
default_factory=dict,
exclude=True,
)
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger, **kwargs):
quart_app = quart.Quart(__name__)
@@ -1062,11 +1066,11 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await self.listeners[type(lb_event)](lb_event, self)
def sync_on_message(event: lark_oapi.im.v1.P2ImMessageReceiveV1):
asyncio.create_task(on_message(event))
self._schedule_inbound_event(on_message(event))
def schedule_on_app_loop(coro):
"""Run a coroutine on the application event loop from sync callbacks."""
return asyncio.run_coroutine_threadsafe(coro, self.ap.event_loop)
return self._schedule_threadsafe_event(coro)
def sync_on_card_action(event):
try:
@@ -1289,7 +1293,13 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
card_form_content={},
card_form_input_defs={},
card_form_inputs={},
card_last_accessed={},
card_cleanup_at=0.0,
card_resume_transitioned=set(),
inbound_event_tasks=set(),
threadsafe_event_futures=set(),
threadsafe_event_lock=threading.Lock(),
tenant_access_tokens={},
seq=1,
listeners={},
quart_app=quart_app,
@@ -1300,6 +1310,45 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
**kwargs,
)
def _schedule_inbound_event(self, coro) -> None:
for task in tuple(self.inbound_event_tasks):
if task.done():
self.inbound_event_tasks.discard(task)
if len(self.inbound_event_tasks) >= self._MAX_INBOUND_EVENTS:
coro.close()
return
task = asyncio.create_task(coro)
self.inbound_event_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
self.inbound_event_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
def _schedule_threadsafe_event(self, coro):
"""Submit one bounded callback from the Lark SDK's sync boundary."""
with self.threadsafe_event_lock:
for future in tuple(self.threadsafe_event_futures):
if future.done():
self.threadsafe_event_futures.discard(future)
if len(self.threadsafe_event_futures) >= self._MAX_INBOUND_EVENTS:
coro.close()
return None
future = asyncio.run_coroutine_threadsafe(coro, self.ap.event_loop)
self.threadsafe_event_futures.add(future)
def done(done_future) -> None:
with self.threadsafe_event_lock:
self.threadsafe_event_futures.discard(done_future)
if not done_future.cancelled():
done_future.exception()
future.add_done_callback(done)
return future
def request_app_ticket(self, api_client, config):
app_id = config['app_id']
app_secret = config['app_secret']
@@ -1376,6 +1425,12 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
'token': tenant_access_token,
'expire_at': int(time.time()) + expire - 300,
}
now = int(time.time())
for cached_key, cached_token in tuple(self.tenant_access_tokens.items()):
if int(cached_token.get('expire_at', 0)) <= now:
self.tenant_access_tokens.pop(cached_key, None)
while len(self.tenant_access_tokens) > self._MAX_TENANT_ACCESS_TOKENS:
self.tenant_access_tokens.pop(next(iter(self.tenant_access_tokens)), None)
def get_tenant_access_token(self, tenant_key: str):
if tenant_key is None or 'isv' != self.config.get('app_type', 'self'):
@@ -1558,6 +1613,8 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
user_msg_id = query.message_event.message_chain.message_id
if user_msg_id:
self.pending_monitoring_msg[user_msg_id] = monitoring_message_id
while len(self.pending_monitoring_msg) > CARD_ID_CACHE_SIZE:
self.pending_monitoring_msg.pop(next(iter(self.pending_monitoring_msg)), None)
except Exception as e:
await self.logger.debug(f'Failed to map message to monitoring message: {e}')
@@ -1570,6 +1627,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
def _next_card_sequence(self, card_id: str, suggested: int = 1) -> int:
"""Return the next strictly increasing sequence for a card update."""
self._touch_card(card_id)
current = self.card_sequence_dict.get(card_id, 0)
next_seq = max(current + 1, suggested)
self.card_sequence_dict[card_id] = next_seq
@@ -1577,6 +1635,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
def _register_card_for_source(self, card_id: str, *source_ids: str) -> None:
"""Register a card_id under one or more source message ids."""
self._touch_card(card_id)
bucket = self.card_id_to_source_ids.setdefault(card_id, set())
for sid in source_ids:
if not sid:
@@ -1596,8 +1655,24 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.card_form_content.pop(card_id, None)
self.card_form_input_defs.pop(card_id, None)
self.card_form_inputs.pop(card_id, None)
self.card_last_accessed.pop(card_id, None)
self.card_resume_transitioned.discard(card_id)
def _touch_card(self, card_id: str) -> None:
now = time.monotonic()
if now - self.card_cleanup_at >= 60 or len(self.card_last_accessed) >= CARD_ID_CACHE_SIZE:
self.card_cleanup_at = now
for stale_card_id, last_accessed in tuple(self.card_last_accessed.items()):
if now - last_accessed >= CARD_ID_CACHE_MAX_LIFETIME:
self._drop_card_state(stale_card_id)
while len(self.card_last_accessed) >= CARD_ID_CACHE_SIZE:
oldest_card_id = min(
self.card_last_accessed,
key=self.card_last_accessed.__getitem__,
)
self._drop_card_state(oldest_card_id)
self.card_last_accessed[card_id] = now
async def create_card_id(self, message_id):
try:
# self.logger.debug('飞书支持stream输出,创建卡片......')
@@ -1793,6 +1868,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.card_id_dict[message_id] = response.data.card_id
card_id = response.data.card_id
self._touch_card(card_id)
self.card_sequence_dict[card_id] = 0
return card_id
@@ -1864,7 +1940,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.reply_to_monitoring_msg[reply_msg_id] = (monitoring_msg_id, time.time())
self._cleanup_monitoring_mapping()
except Exception as e:
asyncio.create_task(self.logger.debug(f'Failed to transfer monitoring mapping in create_message_card: {e}'))
await self.logger.debug(f'Failed to transfer monitoring mapping in create_message_card: {e}')
return True
@@ -2872,8 +2948,8 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
data = await request.json
if 'encrypt' in data:
data = self.cipher.decrypt_string(data['encrypt'])
data = json.loads(data)
encrypted = data['encrypt']
data = await asyncio.to_thread(lambda: json.loads(self.cipher.decrypt_string(encrypted)))
type = self.get_event_type(data)
context = EventContext(data)
if 'url_verification' == type:
@@ -3143,4 +3219,38 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# 所以要设置_auto_reconnect=False,让其不重连。
self.bot._auto_reconnect = False
await self.bot._disconnect()
inbound_tasks = list(self.inbound_event_tasks)
for task in inbound_tasks:
if not task.done():
task.cancel()
if inbound_tasks:
await asyncio.gather(*inbound_tasks, return_exceptions=True)
self.inbound_event_tasks.clear()
with self.threadsafe_event_lock:
threadsafe_futures = list(self.threadsafe_event_futures)
for future in threadsafe_futures:
future.cancel()
if threadsafe_futures:
await asyncio.gather(
*(asyncio.wrap_future(future) for future in threadsafe_futures),
return_exceptions=True,
)
with self.threadsafe_event_lock:
self.threadsafe_event_futures.clear()
self.tenant_access_tokens.clear()
self.card_id_dict.clear()
self.pending_monitoring_msg.clear()
self.reply_to_monitoring_msg.clear()
for card_id in tuple(self.card_last_accessed):
self._drop_card_state(card_id)
self.card_last_accessed.clear()
self.reply_message_card_ids.clear()
self.card_sequence_dict.clear()
self.card_id_to_source_ids.clear()
self.card_streaming_text.clear()
self.card_pre_pause_text.clear()
self.card_form_content.clear()
self.card_form_input_defs.clear()
self.card_form_inputs.clear()
self.card_resume_transitioned.clear()
return False
src/langbot/pkg/platform/sources/legacy/gewechat.py
+13 -8
View File
@@ -6,7 +6,6 @@ import traceback
import time
import re
import copy
import threading
import quart
from langbot.pkg.utils import httpclient
@@ -483,6 +482,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.message_converter = GewechatMessageConverter(config)
self.event_converter = GewechatEventConverter(config)
self.listeners = {}
@self.quart_app.route('/gewechat/callback', methods=['POST'])
async def gewechat_callback():
@@ -518,9 +518,13 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
at_targets = at_targets or []
member_info = []
if at_targets:
member_info = self.bot.get_chatroom_member_detail(self.config['app_id'], target_id, at_targets[::-1])[
'data'
]
member_result = await asyncio.to_thread(
self.bot.get_chatroom_member_detail,
self.config['app_id'],
target_id,
at_targets[::-1],
)
member_info = member_result['data']
# 处理消息组件
for msg in content_list:
@@ -596,7 +600,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
}
if handler := handler_map.get(msg['type']):
handler(msg)
await asyncio.to_thread(handler, msg)
else:
await self.logger.warning(f'未处理的消息类型: {msg["type"]}')
continue
@@ -645,8 +649,9 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
json={'app_id': self.config['app_id']},
) as response:
if response.status != 200:
raise Exception(f'获取gewechat token失败: {await response.text()}')
self.config['token'] = (await response.json())['data']
error = await httpclient.read_text_limited(response)
raise Exception(f'获取gewechat token失败: {error}')
self.config['token'] = (await httpclient.read_json_limited(response))['data']
self.bot = gewechat_client.GewechatClient(f'{self.config["gewechat_url"]}/v2/api', self.config['token'])
@@ -672,7 +677,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
except Exception as e:
raise Exception(f'设置 Gewechat 回调失败， token失效： {e}')
threading.Thread(target=gewechat_login_process).start()
await asyncio.to_thread(gewechat_login_process)
async def shutdown_trigger_placeholder():
while True:
src/langbot/pkg/platform/sources/legacy/nakuru.py
+5 -3
View File
@@ -311,15 +311,17 @@ class NakuruAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
try:
import requests
resp = requests.get(
url='http://{}:{}/get_login_info'.format(self.cfg['host'], self.cfg['http_port']),
resp = await asyncio.to_thread(
requests.get,
'http://{}:{}/get_login_info'.format(self.cfg['host'], self.cfg['http_port']),
headers={'Authorization': 'Bearer ' + self.cfg['token'] if 'token' in self.cfg else ''},
timeout=5,
proxies=None,
)
if resp.status_code == 403:
raise Exception('go-cqhttp拒绝访问，请检查配置文件中nakuru适配器的配置')
self.bot_account_id = int(resp.json()['data']['user_id'])
response_data = await httpclient.parse_json_response(resp)
self.bot_account_id = int(response_data['data']['user_id'])
except Exception:
raise Exception('获取go-cqhttp账号信息失败, 请检查是否已启动go-cqhttp并配置正确')
await self.bot._run()
src/langbot/pkg/platform/sources/legacy/qqbotpy.py
+10 -1
View File
@@ -5,6 +5,7 @@ import typing
import datetime
import re
import traceback
from collections import OrderedDict
import botpy
import botpy.message as botpy_message
@@ -40,7 +41,8 @@ event_handler_mapping = {
}
cached_message_ids = {}
_CACHED_MESSAGE_ID_LIMIT = 10000
cached_message_ids: OrderedDict[str, str] = OrderedDict()
"""由于QQ官方的消息id是字符串，而YiriMirai的消息id是整数，所以需要一个索引来进行转换"""
id_index = 0
@@ -53,6 +55,8 @@ def save_msg_id(message_id: str) -> int:
crt_index = id_index
id_index += 1
cached_message_ids[str(crt_index)] = message_id
while len(cached_message_ids) > _CACHED_MESSAGE_ID_LIMIT:
cached_message_ids.popitem(last=False)
return crt_index
@@ -355,6 +359,7 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.cfg = cfg
self.ap = ap
self.logger = logger
self.cached_official_messages = OrderedDict()
self.group_msg_seq = 1
self.c2c_msg_seq = 1
@@ -490,6 +495,8 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
],
):
self.cached_official_messages[str(message.id)] = message
while len(self.cached_official_messages) > 1000:
self.cached_official_messages.popitem(last=False)
await callback(self.event_converter.target2yiri(message), self)
for event_handler in event_handler_mapping[event_type]:
@@ -519,6 +526,8 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await (await self.bot.start(**self.cfg))
async def kill(self) -> bool:
self.cached_official_messages.clear()
if not self.bot.is_closed():
await self.bot.close()
return True
return True
src/langbot/pkg/platform/sources/line.py
+23 -7
View File
@@ -13,6 +13,7 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
from ..logger import EventLogger
from ...utils import bounded_executor
from linebot.v3 import WebhookHandler
@@ -30,6 +31,14 @@ from linebot.v3.webhooks import (
from linebot.v3.webhook import WebhookParser
from linebot.v3.messaging import MessagingApiBlob
MAX_LINE_MEDIA_BYTES = 10 * 1024 * 1024
def _validate_line_media_content(content: bytes) -> bytes:
if len(content) > MAX_LINE_MEDIA_BYTES:
raise ValueError(f'LINE media exceeds the {MAX_LINE_MEDIA_BYTES}-byte limit')
return content
class LINEMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
@staticmethod
@@ -63,9 +72,13 @@ class LINEMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
elif isinstance(message.message, VideoMessageContent):
pass
elif isinstance(message.message, ImageMessageContent):
message_content = MessagingApiBlob(bot_client).get_message_content(message.message.id)
message_content = await asyncio.to_thread(
MessagingApiBlob(bot_client).get_message_content,
message.message.id,
)
_validate_line_media_content(message_content)
base64_string = base64.b64encode(message_content).decode('utf-8')
base64_string = await asyncio.to_thread(lambda: base64.b64encode(message_content).decode('utf-8'))
# 如果需要Data URI格式（用于直接嵌入HTML等）
# 首先需要知道图片类型，LINE图片通常是JPEG
@@ -173,20 +186,22 @@ class LINEAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
for content in content_list:
if content['type'] == 'text':
self.bot.reply_message_with_http_info(
await asyncio.to_thread(
self.bot.reply_message_with_http_info,
ReplyMessageRequest(
reply_token=message_source.source_platform_object.reply_token,
messages=[TextMessage(text=content['content'])],
)
),
)
elif content['type'] == 'image':
# LINE ImageMessage requires original_content_url and preview_image_url
image_url = content['image']
self.bot.reply_message_with_http_info(
await asyncio.to_thread(
self.bot.reply_message_with_http_info,
ReplyMessageRequest(
reply_token=message_source.source_platform_object.reply_token,
messages=[ImageMessage(original_content_url=image_url, preview_image_url=image_url)],
)
),
)
async def is_muted(self, group_id: int) -> bool:
@@ -266,4 +281,5 @@ class LINEAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await keep_alive()
async def kill(self) -> bool:
pass
await bounded_executor.run_blocking_cleanup(self.api_client.close)
return True
src/langbot/pkg/platform/sources/matrix.py
+98 -27
View File
@@ -5,6 +5,8 @@ import asyncio
import traceback
import base64
import json
import os
from urllib.parse import urlparse
import nio
@@ -16,6 +18,58 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
_MAX_MATRIX_MEDIA_BYTES = 10 * 1024 * 1024
def _decode_matrix_base64_limited(value: str) -> bytes:
if ';base64,' in value:
value = value.split(';base64,', 1)[1]
max_encoded_bytes = 4 * ((_MAX_MATRIX_MEDIA_BYTES + 2) // 3)
if len(value) > max_encoded_bytes:
raise ValueError('Matrix media exceeds the size limit')
decoded = base64.b64decode(value)
if len(decoded) > _MAX_MATRIX_MEDIA_BYTES:
raise ValueError('Matrix media exceeds the size limit')
return decoded
def _read_matrix_file_limited(path: str) -> bytes:
if os.path.getsize(path) > _MAX_MATRIX_MEDIA_BYTES:
raise ValueError('Matrix media exceeds the size limit')
with open(path, 'rb') as file:
body = file.read(_MAX_MATRIX_MEDIA_BYTES + 1)
if len(body) > _MAX_MATRIX_MEDIA_BYTES:
raise ValueError('Matrix media exceeds the size limit')
return body
async def _download_matrix_media_limited(
client: nio.AsyncClient,
mxc_url: str,
) -> tuple[bytes, str]:
parsed = urlparse(mxc_url)
if parsed.scheme != 'mxc' or not parsed.netloc or not parsed.path.strip('/'):
raise ValueError('Invalid Matrix media URL')
method, path = nio.Api.download(
parsed.netloc,
parsed.path.replace('/', ''),
access_token=None,
)
headers = {}
if client.access_token:
headers['Authorization'] = f'Bearer {client.access_token}'
response = await client.send(method, path, headers=headers, timeout=30)
try:
response.raise_for_status()
body = await httpclient.read_limited(
response,
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
)
return body, response.headers.get('Content-Type', 'application/octet-stream')
finally:
response.release()
class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
@staticmethod
async def yiri2target(message_chain: platform_message.MessageChain, client: nio.AsyncClient) -> list[dict]:
@@ -26,17 +80,22 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
elif isinstance(component, platform_message.Image):
image_bytes = None
if component.base64:
b64_data = component.base64
if ';base64,' in b64_data:
b64_data = b64_data.split(';base64,', 1)[1]
image_bytes = base64.b64decode(b64_data)
image_bytes = await asyncio.to_thread(
_decode_matrix_base64_limited,
component.base64,
)
elif component.url:
session = httpclient.get_session()
async with session.get(component.url) as response:
image_bytes = await response.read()
image_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
)
elif component.path:
with open(component.path, 'rb') as f:
image_bytes = f.read()
image_bytes = await asyncio.to_thread(
_read_matrix_file_limited,
str(component.path),
)
if image_bytes:
resp = await client.upload(image_bytes, content_type='image/png')
if isinstance(resp, nio.UploadResponse):
@@ -44,17 +103,22 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
elif isinstance(component, platform_message.File):
file_bytes = None
if component.base64:
b64_data = component.base64
if ';base64,' in b64_data:
b64_data = b64_data.split(';base64,', 1)[1]
file_bytes = base64.b64decode(b64_data)
file_bytes = await asyncio.to_thread(
_decode_matrix_base64_limited,
component.base64,
)
elif component.url:
session = httpclient.get_session()
async with session.get(component.url) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
)
elif component.path:
with open(component.path, 'rb') as f:
file_bytes = f.read()
file_bytes = await asyncio.to_thread(
_read_matrix_file_limited,
str(component.path),
)
if file_bytes:
file_name = getattr(component, 'name', None) or 'file'
resp = await client.upload(file_bytes, content_type='application/octet-stream', filename=file_name)
@@ -86,11 +150,12 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
elif isinstance(event, nio.RoomMessageImage):
mxc_url = event.url
if mxc_url:
resp = await client.download(mxc_url)
if isinstance(resp, nio.DownloadResponse):
b64 = base64.b64encode(resp.body).decode('utf-8')
content_type = resp.content_type or 'image/png'
message_components.append(platform_message.Image(base64=f'data:{content_type};base64,{b64}'))
body, content_type = await _download_matrix_media_limited(
client,
mxc_url,
)
b64 = (await asyncio.to_thread(base64.b64encode, body)).decode('utf-8')
message_components.append(platform_message.Image(base64=f'data:{content_type};base64,{b64}'))
if event.body:
message_components.append(platform_message.Plain(text=event.body))
@@ -431,14 +496,15 @@ class MatrixAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
if not mxc_url:
return
try:
resp = await self.client.download(mxc_url)
if isinstance(resp, nio.DownloadResponse):
b64 = base64.b64encode(resp.body).decode('utf-8')
content_type = resp.content_type or 'image/png'
await self.logger.info(
f'[{_b.user_id}] Bridge 发送了二维码，请扫码登录:',
images=[platform_message.Image(base64=f'data:{content_type};base64,{b64}')],
)
body, content_type = await _download_matrix_media_limited(
self.client,
mxc_url,
)
b64 = (await asyncio.to_thread(base64.b64encode, body)).decode('utf-8')
await self.logger.info(
f'[{_b.user_id}] Bridge 发送了二维码，请扫码登录:',
images=[platform_message.Image(base64=f'data:{content_type};base64,{b64}')],
)
except Exception:
await self.logger.error(
f'[{_b.user_id}] Failed to download bridge QR image: {traceback.format_exc()}'
@@ -672,11 +738,16 @@ class MatrixAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async def kill(self) -> bool:
self._running = False
bridge_tasks = []
for bridge in self._bridges:
if bridge.login_task and not bridge.login_task.done():
bridge.login_task.cancel()
bridge_tasks.append(bridge.login_task)
if bridge.check_task and not bridge.check_task.done():
bridge.check_task.cancel()
bridge_tasks.append(bridge.check_task)
if bridge_tasks:
await asyncio.gather(*bridge_tasks, return_exceptions=True)
if self.client:
await self.client.close()
await self.logger.debug('Matrix adapter stopped')
src/langbot/pkg/platform/sources/officialaccount.py
+1
View File
@@ -164,6 +164,7 @@ class OfficialAccountAdapter(abstract_platform_adapter.AbstractMessagePlatformAd
await keep_alive()
async def kill(self) -> bool:
self.bot.clear()
return False
async def unregister_listener(
src/langbot/pkg/platform/sources/openclaw_weixin.py
+67 -16
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import asyncio
import base64
import os
import traceback
import typing
@@ -26,6 +27,7 @@ from langbot.libs.openclaw_weixin_api.types import (
WeixinMessage,
)
from langbot.pkg.entity.persistence import bot as persistence_bot
from langbot.pkg.utils import httpclient
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
@@ -33,6 +35,10 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.message as platform_message
from langbot.pkg.api.http.context import ExecutionContext
_MAX_OPENCLAW_COMPONENT_BYTES = 10 * 1024 * 1024
class OpenClawWeixinMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
"""Converts between LangBot MessageChain and OpenClaw WeChat message items."""
@@ -112,7 +118,12 @@ class OpenClawWeixinMessageConverter(abstract_platform_adapter.AbstractMessageCo
elif item.type == MessageItem.IMAGE and item.image_item:
if hasattr(item.image_item, '_downloaded_bytes') and item.image_item._downloaded_bytes:
b64 = base64.b64encode(item.image_item._downloaded_bytes).decode('utf-8')
b64 = (
await asyncio.to_thread(
base64.b64encode,
item.image_item._downloaded_bytes,
)
).decode('utf-8')
components.append(platform_message.Image(base64=f'data:image/jpeg;base64,{b64}'))
else:
components.append(platform_message.Unknown(text='[Image]'))
@@ -278,11 +289,36 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
return
try:
ap = self.logger.ap
await ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_bot.Bot)
.where(persistence_bot.Bot.uuid == self._bot_uuid)
.values(adapter_config=self.config)
)
execution_context = getattr(self.logger, 'execution_context', None)
if not isinstance(execution_context, ExecutionContext):
raise RuntimeError('Weixin Bot config persistence requires an ExecutionContext')
if execution_context.bot_uuid != self._bot_uuid:
raise RuntimeError('Weixin Bot UUID does not match its ExecutionContext')
async def persist() -> None:
binding = await ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
if binding.instance_uuid != execution_context.instance_uuid:
raise RuntimeError('Weixin Bot Workspace belongs to another LangBot instance')
await ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_bot.Bot)
.where(persistence_bot.Bot.workspace_uuid == execution_context.workspace_uuid)
.where(persistence_bot.Bot.uuid == self._bot_uuid)
.values(adapter_config=self.config)
)
cloud_runtime = getattr(getattr(ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
tenant_uow = getattr(ap.persistence_mgr, 'tenant_uow', None)
if not callable(tenant_uow):
raise RuntimeError('Cloud adapter persistence requires an explicit tenant UoW')
async with tenant_uow(execution_context.workspace_uuid):
await persist()
else:
await persist()
except Exception as e:
await self.logger.warning(f'Failed to persist adapter config: {e}')
@@ -374,19 +410,30 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
path_val = getattr(component, 'path', None)
if b64_val:
return base64.b64decode(b64_val)
max_encoded_chars = 4 * ((_MAX_OPENCLAW_COMPONENT_BYTES + 2) // 3) + 4
if len(b64_val) > max_encoded_chars:
raise ValueError('OpenClaw media exceeds the size limit')
data = await asyncio.to_thread(base64.b64decode, b64_val)
if len(data) > _MAX_OPENCLAW_COMPONENT_BYTES:
raise ValueError('OpenClaw media exceeds the size limit')
return data
elif url_val and url_val.startswith(('http://', 'https://')):
import aiohttp
async with aiohttp.ClientSession() as session:
async with session.get(url_val) as resp:
if resp.status == 200:
return await resp.read()
session = httpclient.get_session()
async with session.get(url_val) as resp:
if resp.status == 200:
return await httpclient.read_limited(resp)
elif path_val:
import asyncio
if await asyncio.to_thread(os.path.getsize, path_val) > _MAX_OPENCLAW_COMPONENT_BYTES:
raise ValueError('OpenClaw media exceeds the size limit')
with open(path_val, 'rb') as f:
return await asyncio.to_thread(f.read)
def read_file() -> bytes:
with open(path_val, 'rb') as file:
return file.read(_MAX_OPENCLAW_COMPONENT_BYTES + 1)
data = await asyncio.to_thread(read_file)
if len(data) > _MAX_OPENCLAW_COMPONENT_BYTES:
raise ValueError('OpenClaw media exceeds the size limit')
return data
return None
def register_listener(
@@ -517,6 +564,8 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
"""Process a single inbound message from getUpdates."""
if msg.context_token and msg.from_user_id:
self._context_tokens[msg.from_user_id] = msg.context_token
while len(self._context_tokens) > 4096:
self._context_tokens.pop(next(iter(self._context_tokens)), None)
# Download CDN media (files, images) before converting to LangBot events
await self._download_media_items(msg)
@@ -572,6 +621,8 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
await self._poll_task
except asyncio.CancelledError:
pass
self._poll_task = None
self._context_tokens.clear()
await self.client.close()
await self.logger.info('OpenClaw WeChat adapter stopped')
return True
src/langbot/pkg/platform/sources/qqofficial.py
+61 -4
View File
@@ -241,6 +241,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# per (msg_id|event_id) within 60 min, but each reuse needs a
# fresh ``msg_seq`` — re-sending with msg_seq=1 is silently dedup'd.
self._anchor_msg_seq: dict[str, int] = {}
self._background_tasks: set[asyncio.Task] = set()
# Wire button-click handler so webhook mode catches INTERACTION_CREATE.
# (ws mode is wired separately via on_event in _run_websocket so the
@@ -249,6 +250,30 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
async def _on_interaction(event_data: dict, interaction_id: typing.Optional[str]):
await self._handle_interaction_create(event_data, interaction_id)
def _start_background_task(self, coro) -> bool:
"""Start one bounded adapter-side auxiliary task."""
background_tasks = getattr(self, '_background_tasks', None)
if background_tasks is None:
background_tasks = set()
object.__setattr__(self, '_background_tasks', background_tasks)
for task in tuple(background_tasks):
if task.done():
background_tasks.discard(task)
if len(background_tasks) >= 100:
coro.close()
return False
task = asyncio.create_task(coro)
background_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
background_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
async def reply_message(
self,
message_source: platform_events.MessageEvent,
@@ -449,6 +474,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
pass
async def kill(self) -> bool:
task_set = getattr(self, '_background_tasks', set())
background_tasks = list(task_set)
for task in background_tasks:
if not task.done():
task.cancel()
if background_tasks:
await asyncio.gather(*background_tasks, return_exceptions=True)
task_set.clear()
if self._ws_task:
self._ws_task.cancel()
try:
@@ -456,6 +489,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
except asyncio.CancelledError:
pass
self._ws_task = None
await self.bot.close()
self._pending_forms.clear()
self._session_event_ids.clear()
self._anchor_msg_seq.clear()
self._stream_ctx.clear()
self._stream_ctx_ts.clear()
self._fallback_text.clear()
self._fallback_text_ts.clear()
return True
# --------------- 流式输出 ---------------
@@ -473,6 +514,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
for mid in stale_fb:
self._fallback_text.pop(mid, None)
self._fallback_text_ts.pop(mid, None)
while len(self._stream_ctx) > 1000:
oldest = min(self._stream_ctx_ts, key=self._stream_ctx_ts.__getitem__)
self._stream_ctx.pop(oldest, None)
self._stream_ctx_ts.pop(oldest, None)
while len(self._fallback_text) > 1000:
oldest = min(self._fallback_text_ts, key=self._fallback_text_ts.__getitem__)
self._fallback_text.pop(oldest, None)
self._fallback_text_ts.pop(oldest, None)
if stale_ids or stale_fb:
await self.logger.debug(f'Cleaned up {len(stale_ids)} stream contexts, {len(stale_fb)} fallback texts')
@@ -508,6 +557,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# msg_seq=2 instead of being deduplicated by QQ as another seq=1 send.
if source.d_id:
self._anchor_msg_seq[source.d_id] = max(self._anchor_msg_seq.get(source.d_id, 0), 1)
while len(self._anchor_msg_seq) > 4096:
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
ctx = {
'user_openid': source.user_openid,
@@ -577,7 +628,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# 非流式场景（如群聊不支持流式），累积文本后一次性回复
if chunk_text:
# Chunks carry the latest full snapshot, not a text delta.
self._fallback_text[message_id] = chunk_text
self._fallback_text[message_id] = chunk_text[:200000]
self._fallback_text_ts[message_id] = time.time()
if is_final:
full_text = self._fallback_text.pop(message_id, '')
@@ -590,7 +641,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# 累积文本
if chunk_text:
ctx['accumulated_text'] = chunk_text
ctx['accumulated_text'] = chunk_text[:200000]
# 未启动会话时，等第一个有内容的 chunk 来建立会话
if not ctx['session_started']:
@@ -668,6 +719,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
if used >= self._MAX_REPLIES_PER_ANCHOR:
return None
self._anchor_msg_seq[anchor] = used + 1
while len(self._anchor_msg_seq) > 4096:
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
return used + 1
async def _reply_synthetic(
@@ -791,7 +844,9 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
k for k, v in self._session_event_ids.items() if now - v.get('posted_at', 0) > self._PENDING_FORM_TTL
]
for k in stale_e:
self._session_event_ids.pop(k, None)
stale_event = self._session_event_ids.pop(k, None)
if stale_event:
self._anchor_msg_seq.pop(stale_event.get('event_id'), None)
async def _handle_form_chunk(
self,
@@ -973,7 +1028,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
# ACK uses the interaction id, NOT the ws event id.
interaction_id = event_data.get('id') or ''
if interaction_id:
asyncio.create_task(self.bot.ack_interaction(interaction_id, code=0))
self._start_background_task(self.bot.ack_interaction(interaction_id, code=0))
resolved = (event_data.get('data') or {}).get('resolved') or {}
action_id = str(resolved.get('button_data') or resolved.get('button_id') or '').strip()
@@ -1018,6 +1073,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
}
# New anchor → fresh 5-reply budget.
self._anchor_msg_seq[cached_event_id] = 0
while len(self._anchor_msg_seq) > 4096:
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
if self.ap is not None and not ws_event_id:
self.ap.logger.warning(
'QQ Official: INTERACTION_CREATE lacked ws_event_id; '
src/langbot/pkg/platform/sources/satori.py
+21 -7
View File
@@ -18,6 +18,9 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
from langbot.pkg.utils import httpclient
_MAX_GATEWAY_MESSAGE_BYTES = 1024 * 1024
class SatoriMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
@@ -63,7 +66,15 @@ class SatoriMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
padding = 4 - len(raw_b64) % 4
if padding != 4:
raw_b64 += '=' * padding
image_bytes = base64.b64decode(raw_b64)
max_encoded_chars = 4 * ((10 * 1024 * 1024 + 2) // 3) + 4
if len(raw_b64) > max_encoded_chars:
raise ValueError('Satori image exceeds the 10 MiB limit')
image_bytes = await asyncio.to_thread(
base64.b64decode,
raw_b64,
)
if len(image_bytes) > 10 * 1024 * 1024:
raise ValueError('Satori image exceeds the 10 MiB limit')
uploaded_url = await adapter.upload_image(image_bytes, mime_type)
if uploaded_url:
await adapter.logger.info(f'Satori 图片上传成功: {len(image_bytes)} 字节')
@@ -492,7 +503,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
raise ValueError(f'WebSocket URL必须以ws://或wss://开头: {self.endpoint}')
try:
self.ws = await websockets.connect(self.endpoint)
self.ws = await websockets.connect(self.endpoint, max_size=_MAX_GATEWAY_MESSAGE_BYTES)
await asyncio.sleep(0.1)
await self.send_identify()
@@ -584,7 +595,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async def handle_message(self, message: str):
"""Handle WebSocket message"""
try:
data = json.loads(message)
data = await asyncio.to_thread(json.loads, message)
op = data.get('op')
body = data.get('body', {})
@@ -831,9 +842,10 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
try:
async with self.session.request(method, url, headers=headers, json=data) as response:
if response.status == 200:
return await response.json()
result = await httpclient.read_json_limited(response)
return result if isinstance(result, dict) else None
else:
text = await response.text()
text = await httpclient.read_text_limited(response)
await self.logger.error(f'Satori API 请求失败: {response.status} - {text}')
return None
except Exception as e:
@@ -889,7 +901,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
async with self.session.post(url, headers=headers, data=form_data) as response:
if response.status == 200:
result = await response.json()
result = await httpclient.read_json_limited(response)
# The response should contain the URL of the uploaded file
if isinstance(result, dict) and 'url' in result:
return result['url']
@@ -899,7 +911,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await self.logger.warning(f'Satori 图片上传响应格式未知: {result}')
return None
else:
text = await response.text()
text = await httpclient.read_text_limited(response)
await self.logger.error(f'Satori 图片上传失败: {response.status} - {text}')
return None
except Exception as e:
@@ -911,6 +923,8 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self.running = False
if self.heartbeat_task:
self.heartbeat_task.cancel()
await asyncio.gather(self.heartbeat_task, return_exceptions=True)
self.heartbeat_task = None
if self.ws:
try:
await self.ws.close()
src/langbot/pkg/platform/sources/telegram.py
+72 -25
View File
@@ -10,6 +10,8 @@ import typing
import traceback
import json
import base64
import asyncio
import os
import time
import uuid
import pydantic
@@ -22,6 +24,31 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
_MAX_TELEGRAM_MEDIA_BYTES = 10 * 1024 * 1024
def _decode_telegram_base64_limited(value: str) -> bytes:
if ';base64,' in value:
value = value.split(';base64,', 1)[1]
max_encoded_bytes = 4 * ((_MAX_TELEGRAM_MEDIA_BYTES + 2) // 3)
if len(value) > max_encoded_bytes:
raise ValueError('Telegram media exceeds the size limit')
decoded = base64.b64decode(value)
if len(decoded) > _MAX_TELEGRAM_MEDIA_BYTES:
raise ValueError('Telegram media exceeds the size limit')
return decoded
def _read_telegram_file_limited(path: str) -> bytes:
if os.path.getsize(path) > _MAX_TELEGRAM_MEDIA_BYTES:
raise ValueError('Telegram media exceeds the size limit')
with open(path, 'rb') as file:
body = file.read(_MAX_TELEGRAM_MEDIA_BYTES + 1)
if len(body) > _MAX_TELEGRAM_MEDIA_BYTES:
raise ValueError('Telegram media exceeds the size limit')
return body
def _telegram_select_field_options(form_data: dict) -> tuple[str, list[str]]:
"""Return the active select field and its option values."""
field_name = str(form_data.get('_current_input_field') or '').strip()
@@ -86,35 +113,29 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
if isinstance(component, platform_message.Plain):
components.append({'type': 'text', 'text': component.text})
elif isinstance(component, platform_message.Image):
photo_bytes = None
if component.base64:
photo_bytes = base64.b64decode(component.base64)
elif component.url:
session = httpclient.get_session()
async with session.get(component.url) as response:
photo_bytes = await response.read()
elif component.path:
with open(component.path, 'rb') as f:
photo_bytes = f.read()
photo_bytes, _mime_type = await component.get_bytes()
components.append({'type': 'photo', 'photo': photo_bytes})
elif isinstance(component, platform_message.File):
file_bytes = None
if component.base64:
# Strip data URI prefix if present (e.g. "data:application/pdf;base64,...")
b64_data = component.base64
if ';base64,' in b64_data:
b64_data = b64_data.split(';base64,', 1)[1]
file_bytes = base64.b64decode(b64_data)
file_bytes = await asyncio.to_thread(
_decode_telegram_base64_limited,
component.base64,
)
elif component.url:
session = httpclient.get_session()
async with session.get(component.url) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
)
elif component.path:
with open(component.path, 'rb') as f:
file_bytes = f.read()
file_bytes = await asyncio.to_thread(
_read_telegram_file_limited,
str(component.path),
)
file_name = getattr(component, 'name', None) or 'file'
components.append({'type': 'document', 'document': file_bytes, 'filename': file_name})
@@ -152,13 +173,17 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
file_format = ''
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
)
file_format = 'image/jpeg'
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
message_components.append(
platform_message.Image(
url=file.file_path,
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}',
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
)
)
@@ -172,11 +197,15 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
file_format = message.voice.mime_type or 'audio/ogg'
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
)
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
message_components.append(
platform_message.Voice(
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}',
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
length=message.voice.duration,
)
)
@@ -189,16 +218,22 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
file_name = message.document.file_name or 'document'
file_size = message.document.file_size or 0
file_format = message.document.mime_type or 'application/octet-stream'
if file_size > _MAX_TELEGRAM_MEDIA_BYTES:
raise ValueError('Telegram media exceeds the size limit')
file_bytes = None
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
file_bytes = await response.read()
file_bytes = await httpclient.read_limited(
response,
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
)
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
message_components.append(
platform_message.File(
name=file_name,
size=file_size,
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}',
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
)
)
@@ -264,6 +299,8 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
] = {}
_FORM_ACTION_CACHE_TTL = 30 * 60
_MAX_FORM_ACTION_TITLES = 4096
_MAX_STREAM_STATES = 1000
# callback_data -> (display title, pipeline UUID, expiration time, form group id)
_form_action_titles: typing.Dict[str, tuple[str, str, float, str]] = {}
@@ -286,6 +323,8 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
self._form_action_titles.update(
{callback_data: (title, pipeline_uuid, expires_at, group_id) for callback_data, title in mappings.items()}
)
while len(self._form_action_titles) > self._MAX_FORM_ACTION_TITLES:
self._form_action_titles.pop(next(iter(self._form_action_titles)), None)
def _take_form_action_context(self, callback_data: str, now: float | None = None) -> tuple[str, str] | None:
"""Consume a callback and invalidate every button from the same form."""
@@ -446,6 +485,11 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
bot_account_id='',
listeners={},
)
self._form_action_titles = {}
def _cap_stream_states(self) -> None:
while len(self.msg_stream_id) > self._MAX_STREAM_STATES:
self.msg_stream_id.pop(next(iter(self.msg_stream_id)), None)
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
components = await TelegramMessageConverter.yiri2target(message, self.bot)
@@ -554,6 +598,7 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
args = self._build_message_args(chat_id, 'Thinking...', message_thread_id)
send_msg = await self.bot.send_message(**args)
self.msg_stream_id[message_id] = ('message', send_msg.message_id, False)
self._cap_stream_states()
return True
@@ -845,4 +890,6 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
if self.application.updater:
await self.application.updater.stop()
await self.logger.info('Telegram adapter stopped')
self.msg_stream_id.clear()
self._form_action_titles.clear()
return True
src/langbot/pkg/platform/sources/websocket_adapter.py
+259 -48
View File
@@ -1,7 +1,9 @@
"""WebSocket适配器 - 支持双向通信的IM系统"""
import asyncio
import contextvars
import logging
import time
import typing
from datetime import datetime
@@ -13,9 +15,14 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
from ...core import app
from .websocket_manager import WebSocketConnection, is_valid_session_id, ws_connection_manager
from ...core import entities as core_entities
from .websocket_manager import WebSocketConnection, WebSocketScope, is_valid_session_id, ws_connection_manager
logger = logging.getLogger(__name__)
_current_pipeline_uuid: contextvars.ContextVar[str | None] = contextvars.ContextVar(
'websocket_pipeline_uuid',
default=None,
)
class WebSocketMessage(pydantic.BaseModel):
@@ -40,21 +47,82 @@ class WebSocketSession:
stream_message_indexes: dict[str, dict[str, int]] = {}
"""流式消息索引 {pipeline_uuid: {resp_message_id: message_index}}"""
def __init__(self, id: str):
def __init__(
self,
id: str = '',
*,
max_conversations: int = 200,
max_messages: int = 100,
idle_ttl_seconds: int = 86400,
):
self.id = id
self.message_lists = {}
self.stream_message_indexes = {}
self.message_counters: dict[str, int] = {}
self.last_accessed: dict[str, float] = {}
self.max_conversations = max(int(max_conversations), 1)
self.max_messages = max(int(max_messages), 1)
self.idle_ttl_seconds = max(int(idle_ttl_seconds), 1)
def _prune(self, now: float) -> None:
expired = [
key for key, last_accessed in self.last_accessed.items() if now - last_accessed >= self.idle_ttl_seconds
]
for key in expired:
self.reset(key)
overflow = len(self.message_lists) - self.max_conversations + 1
if overflow <= 0:
return
oldest = sorted(self.last_accessed, key=self.last_accessed.get)
for key in oldest[:overflow]:
self.reset(key)
def get_message_list(self, pipeline_uuid: str) -> list[WebSocketMessage]:
now = time.monotonic()
self._prune(now)
if pipeline_uuid not in self.message_lists:
self.message_lists[pipeline_uuid] = []
self.last_accessed[pipeline_uuid] = now
return self.message_lists[pipeline_uuid]
def get_stream_message_indexes(self, pipeline_uuid: str) -> dict[str, int]:
if pipeline_uuid not in self.stream_message_indexes:
self.stream_message_indexes[pipeline_uuid] = {}
self.last_accessed[pipeline_uuid] = time.monotonic()
return self.stream_message_indexes[pipeline_uuid]
def next_message_id(self, conversation_key: str) -> int:
next_id = self.message_counters.get(conversation_key, 0) + 1
self.message_counters[conversation_key] = next_id
return next_id
def append_message(self, conversation_key: str, message: WebSocketMessage) -> None:
messages = self.get_message_list(conversation_key)
messages.append(message)
overflow = len(messages) - self.max_messages
if overflow <= 0:
return
del messages[:overflow]
indexes = self.stream_message_indexes.get(conversation_key, {})
adjusted_indexes = {
response_id: index - overflow for response_id, index in indexes.items() if index >= overflow
}
indexes.clear()
indexes.update(adjusted_indexes)
def reset(self, conversation_key: str) -> None:
self.message_lists.pop(conversation_key, None)
self.stream_message_indexes.pop(conversation_key, None)
self.message_counters.pop(conversation_key, None)
self.last_accessed.pop(conversation_key, None)
def clear(self) -> None:
self.message_lists.clear()
self.stream_message_indexes.clear()
self.message_counters.clear()
self.last_accessed.clear()
class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
"""WebSocket适配器 - 支持双向实时通信"""
@@ -70,7 +138,14 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
ap: app.Application = pydantic.Field(exclude=True)
# 主动推送消息的队列
outbound_message_queue: asyncio.Queue = pydantic.Field(default_factory=asyncio.Queue, exclude=True)
outbound_message_queue: asyncio.Queue = pydantic.Field(
default_factory=lambda: asyncio.Queue(maxsize=100),
exclude=True,
)
inbound_listener_tasks: set[asyncio.Task] = pydantic.Field(
default_factory=set,
exclude=True,
)
"""后端主动推送消息的队列"""
# 流式输出开关
@@ -84,11 +159,26 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
**kwargs,
)
self.websocket_person_session = WebSocketSession(id='websocketperson')
self.websocket_group_session = WebSocketSession(id='websocketgroup')
application = kwargs.get('ap')
instance_data = getattr(getattr(application, 'instance_config', None), 'data', {})
retention = (
instance_data.get('system', {}).get('websocket_retention', {}) if isinstance(instance_data, dict) else {}
)
session_options = {
'max_conversations': retention.get('max_conversations_per_workspace', 200),
'max_messages': retention.get('max_messages_per_conversation', 100),
'idle_ttl_seconds': retention.get('conversation_idle_ttl_seconds', 86400),
}
self.websocket_person_session = WebSocketSession(id='websocketperson', **session_options)
self.websocket_group_session = WebSocketSession(id='websocketgroup', **session_options)
self.bot_account_id = 'websocketbot'
self.outbound_message_queue = asyncio.Queue()
try:
outbound_queue_size = max(int(retention.get('send_queue_size', 100)), 1)
except (TypeError, ValueError):
outbound_queue_size = 100
self.outbound_message_queue = asyncio.Queue(maxsize=outbound_queue_size)
self.inbound_listener_tasks = set()
self.stream_enabled = True
@staticmethod
@@ -113,9 +203,38 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
return None
return pipeline_uuid, session_id
@classmethod
async def _get_connection_from_target(cls, target_id: str):
def _scope(self) -> WebSocketScope:
"""Return this adapter's immutable runtime placement."""
return WebSocketScope.from_context(self.logger.execution_context)
def get_pipeline_uuid_override(self) -> str | None:
"""Return the connection pipeline propagated into the listener task."""
return _current_pipeline_uuid.get()
def _listener_task_done(self, task: asyncio.Task) -> None:
listener_tasks = getattr(self, 'inbound_listener_tasks', None)
if listener_tasks is not None:
listener_tasks.discard(task)
if not task.cancelled():
task.exception()
@staticmethod
def _history_message_chain(message_chain: list[dict]) -> list[dict]:
"""Remove large transient payloads before retaining browser history."""
history = []
for component in message_chain:
copied = dict(component)
if copied.get('base64'):
copied['base64'] = ''
history.append(copied)
return history
async def _get_connection_from_target(self, target_id: str):
"""Resolve a person or group WebSocket launcher to its connection."""
scope = self._scope()
target_value = str(target_id)
for prefix in ('websocket_', 'websocketgroup_'):
if target_value.startswith(prefix):
@@ -123,14 +242,18 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
break
else:
return None
connection = await ws_connection_manager.get_connection(target)
connection = await ws_connection_manager.get_connection(target, scope=scope)
if connection is not None:
return connection
embed_target = cls._parse_embed_target(target_id)
embed_target = self._parse_embed_target(target_id)
if embed_target is not None:
pipeline_uuid, session_id = embed_target
return await ws_connection_manager.get_connection_by_session_id(session_id, pipeline_uuid)
return await ws_connection_manager.get_connection_by_session_id(target)
return await ws_connection_manager.get_connection_by_session_id(
session_id,
scope=scope,
pipeline_uuid=pipeline_uuid,
)
return await ws_connection_manager.get_connection_by_session_id(target, scope=scope)
async def _get_message_context(self, message_source) -> tuple[str, str | None]:
"""Resolve the originating pipeline and browser session for a reply."""
@@ -142,7 +265,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
embed_target = self._parse_embed_target(sender_id)
if embed_target is not None:
return embed_target
return typing.cast(str, self.ap.platform_mgr.websocket_proxy_bot.bot_entity.use_pipeline_uuid), None
raise ValueError('WebSocket reply target is not bound to this adapter scope')
async def send_message(
self,
@@ -160,22 +283,23 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
if connection is not None:
pipeline_uuid = connection.pipeline_uuid
session_id = connection.session_id
scope = connection.scope
else:
embed_target = self._parse_embed_target(target_id)
if embed_target is not None:
pipeline_uuid, session_id = embed_target
else:
pipeline_uuid = typing.cast(
str,
self.ap.platform_mgr.websocket_proxy_bot.bot_entity.use_pipeline_uuid,
)
pipeline_uuid = str(target_id).strip()
if not pipeline_uuid:
raise ValueError('WebSocket target pipeline is required')
session_id = None
scope = self._scope()
session_type = 'group' if target_type == 'group' else 'person'
conversation_key = self._conversation_key(pipeline_uuid, session_id)
session = self.websocket_group_session if session_type == 'group' else self.websocket_person_session
msg_id = len(session.get_message_list(conversation_key)) + 1
msg_id = session.next_message_id(conversation_key)
message_data = WebSocketMessage(
id=msg_id,
@@ -186,7 +310,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
is_final=True,
)
session.get_message_list(conversation_key).append(message_data)
session.append_message(conversation_key, message_data)
await ws_connection_manager.broadcast_to_pipeline(
pipeline_uuid,
@@ -195,6 +319,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
'session_type': session_type,
'data': message_data.model_dump(),
},
scope=scope,
session_type=session_type,
session_id=session_id,
)
@@ -216,10 +341,11 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
)
pipeline_uuid, session_id = await self._get_message_context(message_source)
scope = self._scope()
session_type = 'group' if isinstance(message_source, platform_events.GroupMessage) else 'person'
conversation_key = self._conversation_key(pipeline_uuid, session_id)
msg_id = len(session.get_message_list(conversation_key)) + 1
msg_id = session.next_message_id(conversation_key)
message_data = WebSocketMessage(
id=msg_id,
@@ -230,7 +356,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
is_final=True,
)
session.get_message_list(conversation_key).append(message_data)
session.append_message(conversation_key, message_data)
await ws_connection_manager.broadcast_to_pipeline(
pipeline_uuid,
@@ -239,6 +365,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
'session_type': session_type,
'data': message_data.model_dump(),
},
scope=scope,
session_type=session_type,
session_id=session_id,
)
@@ -262,6 +389,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
)
pipeline_uuid, session_id = await self._get_message_context(message_source)
scope = self._scope()
session_type = 'group' if isinstance(message_source, platform_events.GroupMessage) else 'person'
conversation_key = self._conversation_key(pipeline_uuid, session_id)
message_list = session.get_message_list(conversation_key)
@@ -276,7 +404,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
if existing_index is None or existing_index >= len(message_list):
# 创建新消息
msg_id = len(message_list) + 1
msg_id = session.next_message_id(conversation_key)
message_data = WebSocketMessage(
id=msg_id,
role='assistant',
@@ -287,7 +415,8 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
)
# 立即添加到历史记录（即使is_final=False），以便后续块可以更新它
message_list.append(message_data)
session.append_message(conversation_key, message_data)
message_list = session.get_message_list(conversation_key)
if resp_message_id:
stream_message_indexes[resp_message_id] = len(message_list) - 1
else:
@@ -316,6 +445,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
'session_type': session_type,
'data': message_data.model_dump(),
},
scope=scope,
session_type=session_type,
session_id=session_id,
)
@@ -360,7 +490,11 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
message = await asyncio.wait_for(self.outbound_message_queue.get(), timeout=0.1)
# 广播到所有相关连接
target_id = message.get('target_id', '')
await ws_connection_manager.broadcast_to_pipeline(target_id, message)
await ws_connection_manager.broadcast_to_pipeline(
target_id,
message,
scope=self._scope(),
)
except asyncio.TimeoutError:
pass
@@ -370,9 +504,28 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
async def kill(self):
"""停止适配器"""
pass
await ws_connection_manager.close_scope(self._scope())
listener_tasks = getattr(self, 'inbound_listener_tasks', set())
inbound_tasks = list(listener_tasks)
for task in inbound_tasks:
if not task.done():
task.cancel()
if inbound_tasks:
await asyncio.gather(*inbound_tasks, return_exceptions=True)
listener_tasks.clear()
self.websocket_person_session.clear()
self.websocket_group_session.clear()
while not self.outbound_message_queue.empty():
try:
self.outbound_message_queue.get_nowait()
except asyncio.QueueEmpty:
break
async def _process_image_components(self, message_chain_obj: list):
async def _process_image_components(
self,
connection: WebSocketConnection,
message_chain_obj: list,
):
"""
处理消息链中的图片、语音和文件组件，将 path 转换为 base64
@@ -387,18 +540,36 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
import base64
import mimetypes
storage_mgr = self.ap.storage_mgr
attachments = [
component
for component in message_chain_obj
if component.get('path') and component.get('type') in ('Image', 'Voice', 'File')
]
if not attachments:
return
for component in message_chain_obj:
storage_mgr = self.ap.storage_mgr
execution_context = connection.execution_context
expected_prefix = storage_mgr.scoped_prefix(execution_context, owner_type='upload_image')
for component in attachments:
comp_type = component.get('type', '')
comp_path = component.get('path', '')
if not comp_path or comp_type not in ('Image', 'Voice', 'File'):
continue
if not comp_path.startswith(expected_prefix) or not storage_mgr.is_scoped_object_key(
comp_path,
expected_owner_type='upload_image',
):
await self.logger.warning(f'Rejected {comp_type} attachment outside the WebSocket connection scope')
raise ValueError('Attachment key does not belong to this WebSocket connection')
try:
file_content = await storage_mgr.storage_provider.load(comp_path)
base64_str = base64.b64encode(file_content).decode('utf-8')
file_content = await storage_mgr.load_scoped_object_key(
execution_context,
comp_path,
expected_owner_type='upload_image',
)
base64_str = (await asyncio.to_thread(base64.b64encode, file_content)).decode('utf-8')
lowered = comp_path.lower()
if comp_type == 'Image':
@@ -416,10 +587,15 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
mime_type = mimetypes.guess_type(comp_path)[0] or 'application/octet-stream'
component['base64'] = f'data:{mime_type};base64,{base64_str}'
await storage_mgr.storage_provider.delete(comp_path)
await storage_mgr.delete_scoped_object_key(
execution_context,
comp_path,
expected_owner_type='upload_image',
)
component['path'] = ''
except Exception as e:
await self.logger.error(f'Failed to load {comp_type} file {comp_path}: {e}')
raise
async def handle_websocket_message(
self,
@@ -451,23 +627,23 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
message_chain_obj = message_data.get('message', [])
await self._process_image_components(message_chain_obj)
await self._process_image_components(connection, message_chain_obj)
message_chain = platform_message.MessageChain.model_validate(message_chain_obj)
message_id = len(use_session.get_message_list(conversation_key)) + 1
message_id = use_session.next_message_id(conversation_key)
# 保存用户消息
user_message = WebSocketMessage(
id=message_id,
role='user',
content=str(message_chain),
message_chain=message_chain_obj,
message_chain=self._history_message_chain(message_chain_obj),
timestamp=datetime.now().isoformat(),
connection_id=connection.connection_id,
is_final=True, # 用户消息始终是完整的，非流式
)
use_session.get_message_list(conversation_key).append(user_message)
use_session.append_message(conversation_key, user_message)
await ws_connection_manager.broadcast_to_pipeline(
pipeline_uuid,
@@ -476,6 +652,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
'session_type': session_type,
'data': user_message.model_dump(),
},
scope=connection.scope,
session_type=session_type,
session_id=connection.session_id,
)
@@ -506,11 +683,6 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
sender=sender, message_chain=message_chain, time=datetime.now().timestamp()
)
# 设置流水线UUID (proxy bot always needs it for reply_message routing)
self.ap.platform_mgr.websocket_proxy_bot.bot_entity.use_pipeline_uuid = pipeline_uuid
if owner_bot is not None:
owner_bot.bot_entity.use_pipeline_uuid = pipeline_uuid
# 异步触发事件处理
# Use owner_bot's listeners if available, otherwise fall back to proxy bot
listeners = (
@@ -525,7 +697,38 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
owner_bot.adapter.set_ws_adapter(self)
callback_adapter = owner_bot.adapter if (owner_bot and hasattr(owner_bot, 'adapter')) else self
if event.__class__ in listeners:
asyncio.create_task(listeners[event.__class__](event, callback_adapter))
listener_tasks = getattr(self, 'inbound_listener_tasks', None)
if listener_tasks is None:
listener_tasks = set()
object.__setattr__(self, 'inbound_listener_tasks', listener_tasks)
for task in tuple(listener_tasks):
if task.done():
listener_tasks.discard(task)
if len(listener_tasks) >= 100:
await self.logger.warning('WebSocket inbound listener capacity reached; dropping message')
return
token = _current_pipeline_uuid.set(pipeline_uuid)
try:
task_manager = getattr(self.ap, 'task_mgr', None)
if task_manager is None or not isinstance(getattr(task_manager, 'tasks', None), list):
listener_task = asyncio.create_task(listeners[event.__class__](event, callback_adapter))
else:
listener_task = task_manager.create_task(
listeners[event.__class__](event, callback_adapter),
kind='websocket-message',
name=f'websocket-message-{connection.connection_id}',
scopes=[
core_entities.LifecycleControlScope.APPLICATION,
core_entities.LifecycleControlScope.PLATFORM,
],
instance_uuid=connection.instance_uuid,
workspace_uuid=connection.workspace_uuid,
placement_generation=connection.placement_generation,
).task
listener_tasks.add(listener_task)
listener_task.add_done_callback(self._listener_task_done)
finally:
_current_pipeline_uuid.reset(token)
def get_websocket_messages(
self,
@@ -547,10 +750,14 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
"""Reset one pipeline/client conversation."""
conversation_key = self._conversation_key(pipeline_uuid, session_id)
session = self.websocket_person_session if session_type == 'person' else self.websocket_group_session
if conversation_key in session.message_lists:
session.message_lists[conversation_key] = []
if conversation_key in session.stream_message_indexes:
session.stream_message_indexes[conversation_key] = {}
if isinstance(session, WebSocketSession):
session.reset(conversation_key)
else:
# Compatibility for lightweight adapter doubles.
if conversation_key in session.message_lists:
session.message_lists[conversation_key] = []
if conversation_key in session.stream_message_indexes:
session.stream_message_indexes[conversation_key] = {}
if session_id:
launcher_id = (
@@ -558,11 +765,15 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
if session_type == 'group'
else f'websocket_{pipeline_uuid}:{session_id}'
)
scope = self._scope()
self.ap.sess_mgr.session_list = [
candidate_session
for candidate_session in self.ap.sess_mgr.session_list
if not (
str(
getattr(candidate_session, 'instance_uuid', None) == scope.instance_uuid
and getattr(candidate_session, 'workspace_uuid', None) == scope.workspace_uuid
and getattr(candidate_session, 'placement_generation', None) == scope.placement_generation
and str(
candidate_session.launcher_type.value
if hasattr(candidate_session.launcher_type, 'value')
else candidate_session.launcher_type
src/langbot/pkg/platform/sources/websocket_manager.py
+221 -38
View File
@@ -1,6 +1,7 @@
"""WebSocket连接管理器 - 管理多个并发WebSocket连接"""
import asyncio
import dataclasses
import logging
import typing
import uuid
@@ -8,8 +9,34 @@ from datetime import datetime
import pydantic
from ...api.http.context import ExecutionContext
logger = logging.getLogger(__name__)
_SESSION_FILTER_UNSET = object()
_DEFAULT_SEND_QUEUE_SIZE = 100
@dataclasses.dataclass(frozen=True, slots=True)
class WebSocketScope:
"""Trusted runtime placement carried by every WebSocket connection."""
instance_uuid: str
workspace_uuid: str
placement_generation: int
def __post_init__(self) -> None:
if not self.instance_uuid.strip() or not self.workspace_uuid.strip():
raise ValueError('WebSocket scope requires an instance and Workspace')
if self.placement_generation <= 0:
raise ValueError('WebSocket scope requires a positive placement generation')
@classmethod
def from_context(cls, context: typing.Any) -> 'WebSocketScope':
return cls(
instance_uuid=str(getattr(context, 'instance_uuid', '')),
workspace_uuid=str(getattr(context, 'workspace_uuid', '')),
placement_generation=int(getattr(context, 'placement_generation', 0)),
)
def is_valid_session_id(value: str) -> bool:
@@ -29,6 +56,15 @@ class WebSocketConnection(pydantic.BaseModel):
connection_id: str = pydantic.Field(default_factory=lambda: str(uuid.uuid4()))
"""连接唯一ID"""
instance_uuid: str
"""Owning LangBot instance."""
workspace_uuid: str
"""Owning Workspace."""
placement_generation: int
"""Workspace placement generation captured at connect time."""
pipeline_uuid: str
"""关联的流水线UUID"""
@@ -47,7 +83,10 @@ class WebSocketConnection(pydantic.BaseModel):
last_active: datetime = pydantic.Field(default_factory=datetime.now)
"""最后活跃时间"""
send_queue: asyncio.Queue = pydantic.Field(default_factory=asyncio.Queue, exclude=True)
send_queue: asyncio.Queue = pydantic.Field(
default_factory=lambda: asyncio.Queue(maxsize=_DEFAULT_SEND_QUEUE_SIZE),
exclude=True,
)
"""发送消息队列"""
is_active: bool = True
@@ -56,6 +95,25 @@ class WebSocketConnection(pydantic.BaseModel):
metadata: dict = pydantic.Field(default_factory=dict)
"""连接元数据（可存储额外信息）"""
@property
def scope(self) -> WebSocketScope:
return WebSocketScope(
instance_uuid=self.instance_uuid,
workspace_uuid=self.workspace_uuid,
placement_generation=self.placement_generation,
)
@property
def execution_context(self) -> ExecutionContext:
"""Return the storage/runtime context captured for this connection."""
return ExecutionContext(
instance_uuid=self.instance_uuid,
workspace_uuid=self.workspace_uuid,
placement_generation=self.placement_generation,
pipeline_uuid=self.pipeline_uuid,
)
class WebSocketConnectionManager:
"""WebSocket连接管理器 - 支持多连接并发"""
@@ -64,11 +122,11 @@ class WebSocketConnectionManager:
self.connections: dict[str, WebSocketConnection] = {}
"""所有活跃连接 {connection_id: connection}"""
self.pipeline_connections: dict[str, set[str]] = {}
"""流水线到连接的映射 {pipeline_uuid: {connection_id, ...}}"""
self.pipeline_connections: dict[tuple[str, str, int, str], set[str]] = {}
"""Scoped pipeline to connection mapping."""
self.session_connections: dict[str, set[str]] = {}
"""会话类型到连接的映射 {session_type: {connection_id, ...}}"""
self.session_connections: dict[tuple[str, str, int, str], set[str]] = {}
"""Scoped session-type to connection mapping."""
self._lock = asyncio.Lock()
"""线程锁，保护并发访问"""
@@ -76,40 +134,96 @@ class WebSocketConnectionManager:
async def add_connection(
self,
websocket: typing.Any,
scope: WebSocketScope,
pipeline_uuid: str,
session_type: str,
metadata: dict | None = None,
session_id: str | None = None,
send_queue_size: int = _DEFAULT_SEND_QUEUE_SIZE,
max_connections: int = 1024,
max_connections_per_workspace: int = 32,
) -> WebSocketConnection:
"""Register a WebSocket connection and its optional embed session."""
try:
send_queue_size = max(int(send_queue_size), 1)
except (TypeError, ValueError):
send_queue_size = _DEFAULT_SEND_QUEUE_SIZE
max_connections = max(int(max_connections), 1)
max_connections_per_workspace = max(
min(int(max_connections_per_workspace), max_connections),
1,
)
async with self._lock:
if len(self.connections) >= max_connections:
raise RuntimeError(f'WebSocket connection capacity reached ({max_connections})')
workspace_connection_count = sum(
1
for connection in self.connections.values()
if connection.instance_uuid == scope.instance_uuid
and connection.workspace_uuid == scope.workspace_uuid
and connection.placement_generation == scope.placement_generation
)
if workspace_connection_count >= max_connections_per_workspace:
raise RuntimeError(f'Workspace WebSocket connection capacity reached ({max_connections_per_workspace})')
connection = WebSocketConnection(
instance_uuid=scope.instance_uuid,
workspace_uuid=scope.workspace_uuid,
placement_generation=scope.placement_generation,
pipeline_uuid=pipeline_uuid,
session_type=session_type,
session_id=session_id,
websocket=websocket,
metadata=metadata or {},
send_queue=asyncio.Queue(maxsize=send_queue_size),
)
self.connections[connection.connection_id] = connection
# 更新流水线映射
if pipeline_uuid not in self.pipeline_connections:
self.pipeline_connections[pipeline_uuid] = set()
self.pipeline_connections[pipeline_uuid].add(connection.connection_id)
pipeline_key = self._pipeline_key(scope, pipeline_uuid)
if pipeline_key not in self.pipeline_connections:
self.pipeline_connections[pipeline_key] = set()
self.pipeline_connections[pipeline_key].add(connection.connection_id)
# 更新会话类型映射
if session_type not in self.session_connections:
self.session_connections[session_type] = set()
self.session_connections[session_type].add(connection.connection_id)
session_key = self._session_key(scope, session_type)
if session_key not in self.session_connections:
self.session_connections[session_key] = set()
self.session_connections[session_key].add(connection.connection_id)
logger.debug(
f'WebSocket connection established: {connection.connection_id} '
f'(pipeline={pipeline_uuid}, session_type={session_type})'
f'(workspace={scope.workspace_uuid}, generation={scope.placement_generation}, '
f'pipeline={pipeline_uuid}, session_type={session_type})'
)
return connection
async def close_scope(self, scope: WebSocketScope) -> None:
"""Close and forget every live connection for one runtime placement."""
async with self._lock:
connection_ids = [
connection_id for connection_id, connection in self.connections.items() if connection.scope == scope
]
for connection_id in connection_ids:
connection = self.connections.get(connection_id)
if connection is None:
continue
close = getattr(connection.websocket, 'close', None)
if close is not None:
try:
result = close()
if asyncio.iscoroutine(result):
await result
except Exception:
logger.debug(
'Failed to close WebSocket connection %s',
connection_id,
exc_info=True,
)
await self.remove_connection(connection_id)
async def remove_connection(self, connection_id: str):
"""移除WebSocket连接"""
async with self._lock:
@@ -120,54 +234,103 @@ class WebSocketConnectionManager:
connection.is_active = False
# 从流水线映射中移除
if connection.pipeline_uuid in self.pipeline_connections:
self.pipeline_connections[connection.pipeline_uuid].discard(connection_id)
if not self.pipeline_connections[connection.pipeline_uuid]:
del self.pipeline_connections[connection.pipeline_uuid]
pipeline_key = self._pipeline_key(connection.scope, connection.pipeline_uuid)
if pipeline_key in self.pipeline_connections:
self.pipeline_connections[pipeline_key].discard(connection_id)
if not self.pipeline_connections[pipeline_key]:
del self.pipeline_connections[pipeline_key]
# 从会话类型映射中移除
if connection.session_type in self.session_connections:
self.session_connections[connection.session_type].discard(connection_id)
if not self.session_connections[connection.session_type]:
del self.session_connections[connection.session_type]
session_key = self._session_key(connection.scope, connection.session_type)
if session_key in self.session_connections:
self.session_connections[session_key].discard(connection_id)
if not self.session_connections[session_key]:
del self.session_connections[session_key]
del self.connections[connection_id]
logger.debug(f'WebSocket connection disconnected: {connection_id}')
async def get_connection(self, connection_id: str) -> WebSocketConnection | None:
"""Get a connection by its transport identifier."""
return self.connections.get(connection_id)
@staticmethod
def _pipeline_key(scope: WebSocketScope, pipeline_uuid: str) -> tuple[str, str, int, str]:
return (
scope.instance_uuid,
scope.workspace_uuid,
scope.placement_generation,
pipeline_uuid,
)
@staticmethod
def _session_key(scope: WebSocketScope, session_type: str) -> tuple[str, str, int, str]:
return (
scope.instance_uuid,
scope.workspace_uuid,
scope.placement_generation,
session_type,
)
async def get_connection(
self,
connection_id: str,
*,
scope: WebSocketScope,
) -> WebSocketConnection | None:
"""Get a connection only when it belongs to the expected placement."""
connection = self.connections.get(connection_id)
if connection is None or connection.scope != scope:
return None
return connection
async def get_connection_by_session_id(
self,
session_id: str,
*,
scope: WebSocketScope,
pipeline_uuid: str | None = None,
) -> WebSocketConnection | None:
"""Get an active embed connection by its stable browser session identifier."""
for connection in self.connections.values():
candidates: typing.Iterable[WebSocketConnection]
if pipeline_uuid is not None:
candidates = await self.get_connections_by_pipeline(pipeline_uuid, scope=scope)
else:
candidates = self.connections.values()
for connection in candidates:
if (
connection.session_id == session_id
and connection.is_active
and connection.scope == scope
and (pipeline_uuid is None or connection.pipeline_uuid == pipeline_uuid)
):
return connection
return None
async def get_connections_by_pipeline(self, pipeline_uuid: str) -> list[WebSocketConnection]:
async def get_connections_by_pipeline(
self,
pipeline_uuid: str,
*,
scope: WebSocketScope,
) -> list[WebSocketConnection]:
"""获取指定流水线的所有连接"""
connection_ids = self.pipeline_connections.get(pipeline_uuid, set())
connection_ids = self.pipeline_connections.get(self._pipeline_key(scope, pipeline_uuid), set())
return [self.connections[cid] for cid in connection_ids if cid in self.connections]
async def get_connections_by_session_type(self, session_type: str) -> list[WebSocketConnection]:
async def get_connections_by_session_type(
self,
session_type: str,
*,
scope: WebSocketScope,
) -> list[WebSocketConnection]:
"""获取指定会话类型的所有连接"""
connection_ids = self.session_connections.get(session_type, set())
connection_ids = self.session_connections.get(self._session_key(scope, session_type), set())
return [self.connections[cid] for cid in connection_ids if cid in self.connections]
async def broadcast_to_pipeline(
self,
pipeline_uuid: str,
message: dict,
*,
scope: WebSocketScope,
session_type: str | None = None,
session_id: typing.Any = _SESSION_FILTER_UNSET,
):
@@ -180,7 +343,7 @@ class WebSocketConnectionManager:
session_id: Embed conversation filter. Omit it to broadcast across
conversations; pass ``None`` to target non-embed connections.
"""
connections = await self.get_connections_by_pipeline(pipeline_uuid)
connections = await self.get_connections_by_pipeline(pipeline_uuid, scope=scope)
if session_type is not None:
connections = [conn for conn in connections if conn.session_type == session_type]
@@ -196,13 +359,26 @@ class WebSocketConnectionManager:
async def send_to_connection(self, connection_id: str, message: dict):
"""向指定连接发送消息"""
connection = await self.get_connection(connection_id)
connection = self.connections.get(connection_id)
if not connection or not connection.is_active:
logger.warning(f'Attempt to send message to invalid connection: {connection_id}')
return
try:
await connection.send_queue.put(message)
try:
connection.send_queue.put_nowait(message)
except asyncio.QueueFull:
# A slow or disconnected browser must not backpressure every
# other connection or retain an unbounded response stream.
try:
connection.send_queue.get_nowait()
except asyncio.QueueEmpty:
pass
connection.send_queue.put_nowait(message)
logger.warning(
'WebSocket send queue full; dropped oldest message for connection %s',
connection_id,
)
connection.last_active = datetime.now()
except Exception as e:
logger.error(f'Failed to send message to connection {connection_id}: {e}')
@@ -210,17 +386,24 @@ class WebSocketConnectionManager:
async def update_activity(self, connection_id: str):
"""更新连接活跃时间"""
connection = await self.get_connection(connection_id)
connection = self.connections.get(connection_id)
if connection:
connection.last_active = datetime.now()
def get_stats(self) -> dict:
"""获取连接统计信息"""
def get_stats(self, *, scope: WebSocketScope) -> dict:
"""Return connection statistics for one trusted placement."""
scoped_connections = [connection for connection in self.connections.values() if connection.scope == scope]
pipelines: dict[str, int] = {}
session_types: dict[str, int] = {}
for connection in scoped_connections:
pipelines[connection.pipeline_uuid] = pipelines.get(connection.pipeline_uuid, 0) + 1
session_types[connection.session_type] = session_types.get(connection.session_type, 0) + 1
return {
'total_connections': len(self.connections),
'pipelines': len(self.pipeline_connections),
'connections_by_pipeline': {k: len(v) for k, v in self.pipeline_connections.items()},
'connections_by_session_type': {k: len(v) for k, v in self.session_connections.items()},
'total_connections': len(scoped_connections),
'pipelines': len(pipelines),
'connections_by_pipeline': pipelines,
'connections_by_session_type': session_types,
}
src/langbot/pkg/platform/sources/wechatpad.py
+142 -65
View File
@@ -1,8 +1,6 @@
import requests
import websocket
import json
import time
import httpx
from langbot.libs.wechatpad_api.client import WeChatPadClient
@@ -17,6 +15,7 @@ import threading
import quart
from langbot.pkg.platform.logger import EventLogger
from langbot.pkg.utils import bounded_executor, httpclient
import xml.etree.ElementTree as ET
from typing import Optional, Tuple
from functools import partial
@@ -27,6 +26,8 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
_MAX_GATEWAY_MESSAGE_CHARS = 1024 * 1024
class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger):
@@ -53,12 +54,11 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
content_list.append({'type': 'text', 'content': component.text})
elif isinstance(component, platform_message.Image):
if component.url:
async with httpx.AsyncClient() as client:
response = await client.get(component.url)
if response.status_code == 200:
file_bytes = response.content
base64_str = base64.b64encode(file_bytes).decode('utf-8') # 返回字符串格式
session = httpclient.get_session()
async with session.get(component.url) as response:
if response.status == 200:
file_bytes = await httpclient.read_limited(response)
base64_str = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode('utf-8')
else:
raise Exception('获取文件失败')
# pass
@@ -156,9 +156,19 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
cdnthumburl = img_tag.get('cdnthumburl')
# cdnmidimgurl = img_tag.get('cdnmidimgurl')
image_data = self.bot.cdn_download(aeskey=aeskey, file_type=1, file_url=cdnthumburl)
image_data = await asyncio.to_thread(
self.bot.cdn_download,
aeskey=aeskey,
file_type=1,
file_url=cdnthumburl,
)
if image_data['Data']['FileData'] == '':
image_data = self.bot.cdn_download(aeskey=aeskey, file_type=2, file_url=cdnthumburl)
image_data = await asyncio.to_thread(
self.bot.cdn_download,
aeskey=aeskey,
file_type=2,
file_url=cdnthumburl,
)
base64_str = image_data['Data']['FileData']
# self.logger.info(f"data:image/png;base64,{base64_str}")
@@ -186,7 +196,12 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
if voicemsg is not None:
bufid = voicemsg.get('bufid')
length = voicemsg.get('voicelength')
voice_data = self.bot.get_msg_voice(buf_id=str(bufid), length=int(length), msgid=str(new_msg_id))
voice_data = await asyncio.to_thread(
self.bot.get_msg_voice,
buf_id=str(bufid),
length=int(length),
msgid=str(new_msg_id),
)
audio_base64 = voice_data['Data']['Base64']
# 验证语音数据有效性
@@ -319,7 +334,12 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
# print(aeskey,cdnthumburl)
file_data = self.bot.cdn_download(aeskey=aeskey, file_type=5, file_url=cdnthumburl)
file_data = await asyncio.to_thread(
self.bot.cdn_download,
aeskey=aeskey,
file_type=5,
file_url=cdnthumburl,
)
file_base64 = file_data['Data']['FileData']
# print(file_data)
@@ -538,6 +558,7 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
typing.Type[platform_events.Event],
typing.Callable[[platform_events.Event, abstract_platform_adapter.AbstractMessagePlatformAdapter], None],
] = {}
_MAX_CALLBACK_FUTURES = 100
def __init__(self, config: dict, logger: EventLogger):
quart_app = quart.Quart(__name__)
@@ -556,6 +577,12 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
name='WeChatPad',
bot=bot,
)
self._event_loop: asyncio.AbstractEventLoop | None = None
self._ws_app: websocket.WebSocketApp | None = None
self._ws_thread: threading.Thread | None = None
self._stop_event = threading.Event()
self._callback_futures: set = set()
self._callback_futures_lock = threading.Lock()
async def ws_message(self, data):
"""处理接收到的消息"""
@@ -565,7 +592,7 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
except Exception:
await self.logger.error(f'Error in wechatpad callback: {traceback.format_exc()}')
if event.__class__ in self.listeners:
if event is not None and event.__class__ in self.listeners:
await self.listeners[event.__class__](event, self)
return 'ok'
@@ -580,9 +607,8 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
at_targets = at_targets or []
member_info = []
if at_targets:
member_info = self.bot.get_chatroom_member_detail(
target_id,
)['Data']['member_data']['chatroom_member_list']
member_result = await asyncio.to_thread(self.bot.get_chatroom_member_detail, target_id)
member_info = member_result['Data']['member_data']['chatroom_member_list']
# 处理消息组件
for msg in content_list:
@@ -623,11 +649,35 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
}
if handler := handler_map.get(msg['type']):
handler(msg)
await asyncio.to_thread(handler, msg)
else:
self.logger.warning(f'未处理的消息类型: {msg["type"]}')
await self.logger.warning(f'未处理的消息类型: {msg["type"]}')
continue
def _schedule_ws_message(self, data: dict) -> None:
loop = self._event_loop
if loop is None or loop.is_closed() or self._stop_event.is_set():
return
with self._callback_futures_lock:
if len(self._callback_futures) >= self._MAX_CALLBACK_FUTURES:
return
future = asyncio.run_coroutine_threadsafe(self.ws_message(data), loop)
self._callback_futures.add(future)
def done(completed) -> None:
with self._callback_futures_lock:
self._callback_futures.discard(completed)
if completed.cancelled():
return
try:
completed.result()
except asyncio.CancelledError:
pass
except Exception:
logging.getLogger(__name__).exception('WeChatPad callback failed')
future.add_done_callback(done)
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
"""主动发送消息"""
return await self._handle_message(message, target_id)
@@ -665,86 +715,113 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
pass
async def run_async(self):
self._event_loop = asyncio.get_running_loop()
self._stop_event.clear()
if not self.config['admin_key'] and not self.config['token']:
raise RuntimeError('无wechatpad管理密匙，请填入配置文件后重启')
else:
if self.config['token']:
self.bot = WeChatPadClient(self.config['wechatpad_url'], self.config['token'])
data = self.bot.get_login_status()
data = await asyncio.to_thread(self.bot.get_login_status)
if data['Code'] == 300 and data['Text'] == '你已退出微信':
response = requests.post(
response = await asyncio.to_thread(
requests.post,
f'{self.config["wechatpad_url"]}/admin/GenAuthKey1?key={self.config["admin_key"]}',
json={'Count': 1, 'Days': 365},
timeout=10,
)
if response.status_code != 200:
raise Exception(f'获取token失败: {response.text}')
self.config['token'] = response.json()['Data'][0]
body = await httpclient.response_text(response)
raise Exception(f'获取token失败: {body}')
response_data = await httpclient.parse_json_response(response)
self.config['token'] = response_data['Data'][0]
elif not self.config['token']:
response = requests.post(
response = await asyncio.to_thread(
requests.post,
f'{self.config["wechatpad_url"]}/admin/GenAuthKey1?key={self.config["admin_key"]}',
json={'Count': 1, 'Days': 365},
timeout=10,
)
if response.status_code != 200:
raise Exception(f'获取token失败: {response.text}')
self.config['token'] = response.json()['Data'][0]
body = await httpclient.response_text(response)
raise Exception(f'获取token失败: {body}')
response_data = await httpclient.parse_json_response(response)
self.config['token'] = response_data['Data'][0]
self.bot = WeChatPadClient(self.config['wechatpad_url'], self.config['token'], logger=self.logger)
await self.logger.info(self.config['token'])
thread_1 = threading.Event()
def wechat_login_process():
# 不登录，这些先注释掉，避免登陆态尝试拉qrcode。
# login_data =self.bot.get_login_qr()
# url = login_data['Data']["QrCodeUrl"]
profile = self.bot.get_profile()
# self.logger.info(profile)
self.bot_account_id = profile['Data']['userInfo']['nickName']['str']
self.config['wxid'] = profile['Data']['userInfo']['userName']['str']
thread_1.set()
# asyncio.create_task(wechat_login_process)
threading.Thread(target=wechat_login_process).start()
profile = await asyncio.to_thread(self.bot.get_profile)
self.bot_account_id = profile['Data']['userInfo']['nickName']['str']
self.config['wxid'] = profile['Data']['userInfo']['userName']['str']
def connect_websocket_sync() -> None:
thread_1.wait()
uri = f'{self.config["wechatpad_ws"]}/GetSyncMsg?key={self.config["token"]}'
print(f'Connecting to WebSocket: {uri}')
def on_message(ws, message):
try:
if len(message) > _MAX_GATEWAY_MESSAGE_CHARS:
logging.getLogger(__name__).warning('WeChatPad WebSocket message exceeds the size limit')
return
data = json.loads(message)
# 这里需要确保ws_message是同步的，或者使用asyncio.run调用异步方法
asyncio.run(self.ws_message(data))
self._schedule_ws_message(data)
except json.JSONDecodeError:
self.logger.error(f'Non-JSON message: {message[:100]}...')
logging.getLogger(__name__).warning('WeChatPad received a non-JSON message')
def on_error(ws, error):
self.logger.error(f'WebSocket error: {str(error)[:200]}')
logging.getLogger(__name__).warning('WeChatPad WebSocket error: %s', str(error)[:200])
def on_close(ws, close_status_code, close_msg):
self.logger.info('WebSocket closed, reconnecting...')
time.sleep(5)
connect_websocket_sync() # 自动重连
logging.getLogger(__name__).info('WeChatPad WebSocket closed')
def on_open(ws):
self.logger.info('WebSocket connected successfully!')
logging.getLogger(__name__).info('WeChatPad WebSocket connected')
ws = websocket.WebSocketApp(
uri, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open
)
ws.run_forever(ping_interval=60, ping_timeout=20)
while not self._stop_event.is_set():
ws = websocket.WebSocketApp(
uri,
on_message=on_message,
on_error=on_error,
on_close=on_close,
on_open=on_open,
)
self._ws_app = ws
ws.run_forever(ping_interval=60, ping_timeout=20)
self._ws_app = None
if not self._stop_event.wait(5):
logging.getLogger(__name__).info('Reconnecting WeChatPad WebSocket')
# 直接调用同步版本（会阻塞）
# connect_websocket_sync()
# 这行代码会在WebSocket连接断开后才会执行
thread = threading.Thread(target=connect_websocket_sync, name='WebSocketClientThread', daemon=True)
thread.start()
self.logger.info('WebSocket client thread started')
self._ws_thread = threading.Thread(
target=connect_websocket_sync,
name='WebSocketClientThread',
daemon=True,
)
self._ws_thread.start()
await self.logger.info('WebSocket client thread started')
while not self._stop_event.is_set() and self._ws_thread.is_alive():
await asyncio.sleep(1)
if not self._stop_event.is_set():
raise RuntimeError('WeChatPad WebSocket client thread exited unexpectedly')
async def kill(self) -> bool:
pass
self._stop_event.set()
ws = self._ws_app
if ws is not None:
await bounded_executor.run_blocking_cleanup(ws.close)
with self._callback_futures_lock:
futures = list(self._callback_futures)
for future in futures:
future.cancel()
if futures:
await asyncio.gather(
*(asyncio.wrap_future(future) for future in futures),
return_exceptions=True,
)
thread = self._ws_thread
if thread is not None and thread.is_alive():
await bounded_executor.run_blocking_cleanup(thread.join, 5)
self._ws_thread = None
self._ws_app = None
self._event_loop = None
return True
src/langbot/pkg/platform/sources/wecom.py
+1
View File
@@ -321,6 +321,7 @@ class WecomAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await keep_alive()
async def kill(self) -> bool:
await self.bot.close()
return False
async def unregister_listener(
src/langbot/pkg/platform/sources/wecombot.py
+11
View File
@@ -516,6 +516,8 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
# do work. Lazy-create on first call.
object.__setattr__(self, '_synthetic_buffers', {})
buffers: dict[str, str] = self._synthetic_buffers
if buf_key not in buffers and len(buffers) >= 100:
buffers.pop(next(iter(buffers)), None)
if content and not form_data:
previous = buffers.get(buf_key, '')
if previous and content.startswith(previous):
@@ -524,6 +526,8 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
buffers[buf_key] = previous
else:
buffers[buf_key] = previous + content
if len(buffers[buf_key]) > 200000:
buffers[buf_key] = buffers[buf_key][-200000:]
if not is_final:
return {'stream': True, 'synthetic': True, 'buffered': True}
@@ -613,7 +617,11 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
'chat_id': chat_id,
'stream_id': '',
'req_id': '',
'created_at': time.monotonic(),
}
prune = getattr(self.bot, '_prune_pending_forms', None)
if callable(prune):
prune()
return payload
async def send_message(self, target_type, target_id, message):
@@ -745,10 +753,13 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await keep_alive()
async def kill(self) -> bool:
if hasattr(self, '_synthetic_buffers'):
self._synthetic_buffers.clear()
_ws_mode = not self.config.get('enable-webhook', False)
if _ws_mode:
await self.bot.disconnect()
return True
await self.bot.close()
return False
async def unregister_listener(
src/langbot/pkg/platform/sources/wecomcs.py
+2
View File
@@ -254,6 +254,8 @@ class WecomCSAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
await keep_alive()
async def kill(self) -> bool:
self.bot.clear()
await self.bot.close()
return False
async def is_muted(self, group_id: int) -> bool: