xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-09 04:40:57 +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/libs/coze_server_api/client.py
+42 -8
View File
@@ -6,6 +6,22 @@ from typing import Dict, List, Any, AsyncGenerator
import os
from pathlib import Path
from langbot.pkg.utils import httpclient
_MAX_COZE_RESPONSE_BYTES = 16 * 1024 * 1024
_MAX_COZE_EVENT_BYTES = 1024 * 1024
_MAX_COZE_MEDIA_BYTES = 10 * 1024 * 1024
def _read_local_media_limited(path: Path) -> bytes:
if path.stat().st_size > _MAX_COZE_MEDIA_BYTES:
raise ValueError('Coze upload exceeds the size limit')
with path.open('rb') as handle:
body = handle.read(_MAX_COZE_MEDIA_BYTES + 1)
if len(body) > _MAX_COZE_MEDIA_BYTES:
raise ValueError('Coze upload exceeds the size limit')
return body
class AsyncCozeAPIClient:
def __init__(self, api_key: str, api_base: str = 'https://api.coze.cn'):
@@ -58,19 +74,24 @@ class AsyncCozeAPIClient:
if isinstance(file, Path):
if not file.exists():
raise ValueError(f'File not found: {file}')
with open(file, 'rb') as f:
file = f.read()
file = await asyncio.to_thread(_read_local_media_limited, file)
# 处理文件路径字符串
elif isinstance(file, str):
if not os.path.isfile(file):
raise ValueError(f'File not found: {file}')
with open(file, 'rb') as f:
file = f.read()
file = await asyncio.to_thread(
_read_local_media_limited,
Path(file),
)
# 处理文件对象
elif hasattr(file, 'read'):
file = file.read()
file = await asyncio.to_thread(file.read, _MAX_COZE_MEDIA_BYTES + 1)
if not isinstance(file, (bytes, bytearray)):
raise ValueError('Unsupported Coze upload type')
if len(file) > _MAX_COZE_MEDIA_BYTES:
raise ValueError('Coze upload exceeds the size limit')
session = await self.coze_session()
url = f'{self.api_base}/v1/files/upload'
@@ -87,13 +108,18 @@ class AsyncCozeAPIClient:
if response.status == 401:
raise Exception('Coze API 认证失败，请检查 API Key 是否正确')
response_text = await response.text()
response_text = (
await httpclient.read_limited(
response,
max_bytes=_MAX_COZE_EVENT_BYTES,
)
).decode('utf-8', errors='replace')
if response.status != 200:
raise Exception(f'文件上传失败，状态码: {response.status}, 响应: {response_text}')
try:
result = await response.json()
except json.JSONDecodeError:
result = json.loads(response_text)
except (json.JSONDecodeError, TypeError):
raise Exception(f'文件上传响应解析失败: {response_text}')
if result.get('code') != 0:
@@ -158,7 +184,15 @@ class AsyncCozeAPIClient:
if response.status != 200:
raise Exception(f'Coze API 流式请求失败，状态码: {response.status}')
total_bytes = 0
chunk_type = 'message'
chunk_data = ''
async for chunk in response.content:
total_bytes += len(chunk)
if total_bytes > _MAX_COZE_RESPONSE_BYTES:
raise Exception('Coze API stream exceeds the runtime limit')
if len(chunk) > _MAX_COZE_EVENT_BYTES:
raise Exception('Coze API event exceeds the runtime limit')
chunk = chunk.decode('utf-8')
if chunk != '\n':
if chunk.startswith('event:'):
src/langbot/libs/deerflow_api/client.py
+33 -5
View File
@@ -12,10 +12,26 @@ from collections.abc import AsyncGenerator
import httpx
from langbot.pkg.utils import httpclient
from .errors import DeerFlowAPIError
SSE_MAX_BUFFER_CHARS = 1_048_576
SSE_MAX_TOTAL_BYTES = 16 * 1024 * 1024
ERROR_BODY_MAX_BYTES = 1024 * 1024
async def _read_error_body(response: httpx.Response) -> str:
body = bytearray()
async for chunk in response.aiter_bytes(8192):
body.extend(chunk)
if len(body) > ERROR_BODY_MAX_BYTES:
raise DeerFlowAPIError(
operation='read error response',
body='response exceeds the runtime limit',
)
return body.decode('utf-8', errors='replace')
def _normalize_sse_newlines(text: str) -> str:
@@ -94,6 +110,7 @@ class AsyncDeerFlowClient:
async with httpx.AsyncClient(
trust_env=True,
timeout=timeout,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
response = await client.post(
url,
@@ -101,13 +118,14 @@ class AsyncDeerFlowClient:
json=payload,
)
if response.status_code not in (200, 201):
body = await httpclient.response_text(response)
raise DeerFlowAPIError(
operation='create thread',
status=response.status_code,
body=response.text,
body=body,
url=url,
)
return response.json()
return await httpclient.parse_json_response(response)
async def delete_thread(self, thread_id: str, timeout: float = 20) -> None:
"""删除指定 thread"""
@@ -116,13 +134,15 @@ class AsyncDeerFlowClient:
async with httpx.AsyncClient(
trust_env=True,
timeout=timeout,
event_hooks=httpclient.httpx_response_limit_hooks(),
) as client:
response = await client.delete(url, headers=self.headers)
if response.status_code not in (200, 202, 204, 404):
body = await httpclient.response_text(response)
raise DeerFlowAPIError(
operation='delete thread',
status=response.status_code,
body=response.text,
body=body,
url=url,
thread_id=thread_id,
)
@@ -163,19 +183,27 @@ class AsyncDeerFlowClient:
json=payload,
) as resp:
if resp.status_code != 200:
body = await resp.aread()
raise DeerFlowAPIError(
operation='runs/stream request',
status=resp.status_code,
body=body.decode('utf-8', errors='replace'),
body=await _read_error_body(resp),
url=url,
thread_id=thread_id,
)
decoder = codecs.getincrementaldecoder('utf-8')('replace')
buffer = ''
total_bytes = 0
async for chunk in resp.aiter_bytes(8192):
total_bytes += len(chunk)
if total_bytes > SSE_MAX_TOTAL_BYTES:
raise DeerFlowAPIError(
operation='runs/stream response',
body='response exceeds the runtime limit',
url=url,
thread_id=thread_id,
)
buffer += _normalize_sse_newlines(decoder.decode(chunk))
while '\n\n' in buffer:
src/langbot/libs/dify_service_api/v1/client.py
+176 -111
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import httpx
import typing
import json
@@ -8,6 +9,75 @@ from .errors import DifyAPIError
from pathlib import Path
import os
_MAX_DIFY_RESPONSE_BYTES = 1024 * 1024
_MAX_DIFY_SSE_LINE_BYTES = 1024 * 1024
_MAX_DIFY_STREAM_BYTES = 16 * 1024 * 1024
_MAX_DIFY_UPLOAD_BYTES = 10 * 1024 * 1024
async def _read_limited_response(
response: httpx.Response,
*,
max_bytes: int = _MAX_DIFY_RESPONSE_BYTES,
) -> bytes:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > max_bytes:
raise DifyAPIError(f'Remote response exceeds the {max_bytes}-byte limit')
except (TypeError, ValueError):
pass
body = bytearray()
async for chunk in response.aiter_bytes(chunk_size=8192):
body.extend(chunk)
if len(body) > max_bytes:
raise DifyAPIError(f'Remote response exceeds the {max_bytes}-byte limit')
return bytes(body)
async def _iter_sse_json(
response: httpx.Response,
) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
"""Parse Dify's one-JSON-per-data-line SSE without unbounded line buffering."""
buffer = bytearray()
total = 0
async for chunk in response.aiter_bytes(chunk_size=8192):
total += len(chunk)
if total > _MAX_DIFY_STREAM_BYTES:
raise DifyAPIError('Dify SSE stream exceeds the runtime limit')
buffer.extend(chunk)
while b'\n' in buffer:
raw_line, _, remainder = buffer.partition(b'\n')
buffer = bytearray(remainder)
if len(raw_line) > _MAX_DIFY_SSE_LINE_BYTES:
raise DifyAPIError('Dify SSE event exceeds the runtime limit')
line = raw_line.rstrip(b'\r').strip()
if not line or not line.startswith(b'data:'):
continue
payload = json.loads(line[5:].decode('utf-8', errors='replace'))
if isinstance(payload, dict):
yield payload
if len(buffer) > _MAX_DIFY_SSE_LINE_BYTES:
raise DifyAPIError('Dify SSE event exceeds the runtime limit')
line = bytes(buffer).rstrip(b'\r').strip()
if line.startswith(b'data:'):
payload = json.loads(line[5:].decode('utf-8', errors='replace'))
if isinstance(payload, dict):
yield payload
def _read_local_file_limited(path: Path) -> bytes:
if path.stat().st_size > _MAX_DIFY_UPLOAD_BYTES:
raise ValueError('Dify upload exceeds the size limit')
with path.open('rb') as handle:
body = handle.read(_MAX_DIFY_UPLOAD_BYTES + 1)
if len(body) > _MAX_DIFY_UPLOAD_BYTES:
raise ValueError('Dify upload exceeds the size limit')
return body
class AsyncDifyServiceClient:
"""Dify Service API 客户端"""
@@ -22,6 +92,21 @@ class AsyncDifyServiceClient:
) -> None:
self.api_key = api_key
self.base_url = base_url
self._client: httpx.AsyncClient | None = None
def _get_client(self) -> httpx.AsyncClient:
if self._client is None:
self._client = httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
)
return self._client
async def aclose(self) -> None:
client = self._client
self._client = None
if client is not None:
await client.aclose()
async def chat_messages(
self,
@@ -38,37 +123,32 @@ class AsyncDifyServiceClient:
if response_mode != 'streaming':
raise DifyAPIError('当前仅支持 streaming 模式')
async with httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
timeout=timeout,
) as client:
payload = {
'inputs': inputs,
'query': query,
'user': user,
'response_mode': response_mode,
'conversation_id': conversation_id,
'files': files,
'model_config': model_config or {},
}
client = self._get_client()
payload = {
'inputs': inputs,
'query': query,
'user': user,
'response_mode': response_mode,
'conversation_id': conversation_id,
'files': files,
'model_config': model_config or {},
}
async with client.stream(
'POST',
'/chat-messages',
headers={
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json',
},
json=payload,
) as r:
async for chunk in r.aiter_lines():
if r.status_code != 200:
raise DifyAPIError(f'{r.status_code} {chunk}')
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
yield json.loads(chunk[5:])
async with client.stream(
'POST',
'/chat-messages',
headers={
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json',
},
json=payload,
timeout=timeout,
) as r:
if r.status_code != 200:
body = await _read_limited_response(r)
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
async for event in _iter_sse_json(r):
yield event
async def workflow_run(
self,
@@ -82,32 +162,27 @@ class AsyncDifyServiceClient:
if response_mode != 'streaming':
raise DifyAPIError('当前仅支持 streaming 模式')
async with httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
client = self._get_client()
async with client.stream(
'POST',
'/workflows/run',
headers={
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json',
},
json={
'inputs': inputs,
'user': user,
'response_mode': response_mode,
'files': files,
},
timeout=timeout,
) as client:
async with client.stream(
'POST',
'/workflows/run',
headers={
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json',
},
json={
'inputs': inputs,
'user': user,
'response_mode': response_mode,
'files': files,
},
) as r:
async for chunk in r.aiter_lines():
if r.status_code != 200:
raise DifyAPIError(f'{r.status_code} {chunk}')
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
yield json.loads(chunk[5:])
) as r:
if r.status_code != 200:
body = await _read_limited_response(r)
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
async for event in _iter_sse_json(r):
yield event
async def workflow_submit(
self,
@@ -129,41 +204,38 @@ class AsyncDifyServiceClient:
'Content-Type': 'application/json',
}
async with httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
client = self._get_client()
# Step 1: Submit the form
payload: dict[str, typing.Any] = {
'inputs': inputs if isinstance(inputs, dict) else {},
'user': user,
'action': action,
}
async with client.stream(
'POST',
f'/form/human_input/{form_token}',
headers=headers,
json=payload,
timeout=timeout,
) as client:
# Step 1: Submit the form
payload: dict[str, typing.Any] = {
'inputs': inputs if isinstance(inputs, dict) else {},
'user': user,
'action': action,
}
submit_resp = await client.post(
f'/form/human_input/{form_token}',
headers=headers,
json=payload,
)
) as submit_resp:
submit_body = await _read_limited_response(submit_resp)
if submit_resp.status_code != 200:
raise DifyAPIError(f'{submit_resp.status_code} {submit_resp.text}')
raise DifyAPIError(f'{submit_resp.status_code} {submit_body.decode(errors="replace")}')
# Step 2: Stream resumed workflow events
async with client.stream(
'GET',
f'/workflow/{workflow_run_id}/events',
headers={'Authorization': f'Bearer {self.api_key}'},
params={'user': user},
) as r:
if r.status_code != 200:
body = (await r.aread()).decode(errors='replace')
raise DifyAPIError(f'{r.status_code} {body}')
async for chunk in r.aiter_lines():
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
yield json.loads(chunk[5:])
# Step 2: Stream resumed workflow events
async with client.stream(
'GET',
f'/workflow/{workflow_run_id}/events',
headers={'Authorization': f'Bearer {self.api_key}'},
params={'user': user},
timeout=timeout,
) as r:
if r.status_code != 200:
body = await _read_limited_response(r)
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
async for event in _iter_sse_json(r):
yield event
async def upload_file(
self,
@@ -175,37 +247,30 @@ class AsyncDifyServiceClient:
if isinstance(file, Path):
if not file.exists():
raise ValueError(f'File not found: {file}')
with open(file, 'rb') as f:
file = f.read()
file = await asyncio.to_thread(_read_local_file_limited, file)
# 处理文件路径字符串
elif isinstance(file, str):
if not os.path.isfile(file):
raise ValueError(f'File not found: {file}')
with open(file, 'rb') as f:
file = f.read()
file = await asyncio.to_thread(_read_local_file_limited, Path(file))
# 处理文件对象
elif hasattr(file, 'read'):
file = file.read()
async with httpx.AsyncClient(
base_url=self.base_url,
trust_env=True,
file = await asyncio.to_thread(file.read, _MAX_DIFY_UPLOAD_BYTES + 1)
if len(file) > _MAX_DIFY_UPLOAD_BYTES:
raise ValueError('Dify upload exceeds the size limit')
client = self._get_client()
# multipart/form-data
async with client.stream(
'POST',
'/files/upload',
headers={'Authorization': f'Bearer {self.api_key}'},
files={'file': file},
data={'user': user},
timeout=timeout,
) as client:
# multipart/form-data
response = await client.post(
'/files/upload',
headers={'Authorization': f'Bearer {self.api_key}'},
files={
'file': file,
},
data={
'user': user,
},
)
) as response:
body = await _read_limited_response(response)
if response.status_code != 201:
raise DifyAPIError(f'{response.status_code} {response.text}')
return response.json()
raise DifyAPIError(f'{response.status_code} {body.decode(errors="replace")}')
return json.loads(body)
src/langbot/libs/dingtalk_api/api.py
+133 -46
View File
@@ -7,6 +7,7 @@ import time
import typing
import uuid
import urllib.parse
from contextlib import asynccontextmanager
from typing import Awaitable, Callable, Optional
import dingtalk_stream # type: ignore
import websockets
@@ -15,12 +16,42 @@ from .card_callback import DingTalkCardActionHandler
from .dingtalkevent import DingTalkEvent
import httpx
import traceback
from langbot.pkg.utils import httpclient
_stdout_logger = logging.getLogger('langbot.dingtalk_api')
DINGTALK_OPENAPI_BASE = 'https://api.dingtalk.com'
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
_MAX_GATEWAY_MESSAGE_BYTES = 1024 * 1024
def _read_local_media_limited(file_path: str) -> bytes:
if os.path.getsize(file_path) > _MAX_MEDIA_BYTES:
raise ValueError('DingTalk media exceeds the size limit')
with open(file_path, 'rb') as file:
body = file.read(_MAX_MEDIA_BYTES + 1)
if len(body) > _MAX_MEDIA_BYTES:
raise ValueError('DingTalk media exceeds the size limit')
return body
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > _MAX_MEDIA_BYTES:
raise ValueError('DingTalk media exceeds the size limit')
except (TypeError, ValueError) as exc:
if 'exceeds' in str(exc):
raise
body = bytearray()
async for chunk in response.aiter_bytes():
body.extend(chunk)
if len(body) > _MAX_MEDIA_BYTES:
raise ValueError('DingTalk media exceeds the size limit')
return bytes(body)
def _stringify_card_param_map(card_param_map: Optional[dict]) -> dict:
@@ -44,6 +75,8 @@ def _stringify_card_param_map(card_param_map: Optional[dict]) -> dict:
class DingTalkClient:
_MAX_INBOUND_TASKS = 100
def __init__(
self,
client_id: str,
@@ -86,6 +119,37 @@ class DingTalkClient:
self.legacy_access_token = ''
self.legacy_access_token_expiry_time: typing.Optional[float] = None
self._stopped = False # Flag to control the event loop
self._inbound_tasks: set[asyncio.Task] = set()
self._http_client: httpx.AsyncClient | None = None
@asynccontextmanager
async def _http_client_context(self):
"""Reuse one connection pool while preserving existing call structure."""
if self._http_client is None or self._http_client.is_closed:
self._http_client = httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks())
yield self._http_client
def _start_inbound_task(self, coro: typing.Coroutine) -> bool:
"""Start one bounded inbound callback task."""
for task in tuple(self._inbound_tasks):
if task.done():
self._inbound_tasks.discard(task)
if len(self._inbound_tasks) >= self._MAX_INBOUND_TASKS:
coro.close()
return False
task = asyncio.create_task(coro)
self._inbound_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
self._inbound_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
async def _on_card_action(self, payload: dict) -> None:
"""Dispatch a parsed card-action payload to the adapter callback."""
@@ -101,11 +165,11 @@ class DingTalkClient:
url = 'https://api.dingtalk.com/v1.0/oauth2/accessToken'
headers = {'Content-Type': 'application/json'}
data = {'appKey': self.key, 'appSecret': self.secret}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
try:
response = await client.post(url, json=data, headers=headers)
if response.status_code == 200:
response_data = response.json()
response_data = await httpclient.parse_json_response(response)
self.access_token = response_data.get('accessToken')
expires_in = int(response_data.get('expireIn', 7200))
self.access_token_expiry_time = time.time() + expires_in - 60
@@ -129,28 +193,28 @@ class DingTalkClient:
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
headers = {'x-acs-dingtalk-access-token': self.access_token}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=params)
if response.status_code == 200:
result = response.json()
result = await httpclient.parse_json_response(response)
download_url = result.get('downloadUrl')
else:
await self.logger.error(f'failed to get download url: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'failed to get download url: {error_payload}')
if download_url:
return await self.download_url_to_base64(download_url)
async def download_url_to_base64(self, download_url):
async with httpx.AsyncClient() as client:
response = await client.get(download_url)
if response.status_code == 200:
file_bytes = response.content
mime_type = response.headers.get('Content-Type', 'application/octet-stream')
base64_str = base64.b64encode(file_bytes).decode('utf-8')
return f'data:{mime_type};base64,{base64_str}'
else:
await self.logger.error(f'failed to get files: {response.json()}')
async with self._http_client_context() as client:
async with client.stream('GET', download_url) as response:
if response.status_code == 200:
file_bytes = await _read_httpx_media_limited(response)
mime_type = response.headers.get('Content-Type', 'application/octet-stream')
base64_str = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode('utf-8')
return f'data:{mime_type};base64,{base64_str}'
error_body = await _read_httpx_media_limited(response)
await self.logger.error(f'failed to get files: {error_body[:300]!r}')
async def get_audio_url(self, download_code: str):
if not await self.check_access_token():
@@ -158,17 +222,19 @@ class DingTalkClient:
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
headers = {'x-acs-dingtalk-access-token': self.access_token}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=params)
if response.status_code == 200:
result = response.json()
result = await httpclient.parse_json_response(response)
download_url = result.get('downloadUrl')
if download_url:
return await self.download_url_to_base64(download_url)
else:
await self.logger.error(f'failed to get audio: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'failed to get audio: {error_payload}')
else:
raise Exception(f'Error: {response.status_code}, {response.text}')
body = await httpclient.response_text(response)
raise Exception(f'Error: {response.status_code}, {body}')
async def get_file_url(self, download_code: str):
if not await self.check_access_token():
@@ -176,17 +242,19 @@ class DingTalkClient:
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
headers = {'x-acs-dingtalk-access-token': self.access_token}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=params)
if response.status_code == 200:
result = response.json()
result = await httpclient.parse_json_response(response)
download_url = result.get('downloadUrl')
if download_url:
return download_url
else:
await self.logger.error(f'failed to get file: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'failed to get file: {error_payload}')
else:
raise Exception(f'Error: {response.status_code}, {response.text}')
body = await httpclient.response_text(response)
raise Exception(f'Error: {response.status_code}, {body}')
async def update_incoming_message(self, message):
"""异步更新 DingTalkClient 中的 incoming_message"""
@@ -503,12 +571,13 @@ class DingTalkClient:
len(content),
)
try:
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=data)
response_body = await httpclient.response_text(response, max_chars=500)
_stdout_logger.info(
'DingTalk send_proactive_message_to_one response: status=%d body=%s',
response.status_code,
response.text[:500],
response_body,
)
if response.status_code == 200:
return
@@ -535,7 +604,7 @@ class DingTalkClient:
'msgParam': json.dumps({'content': content}),
}
try:
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=data)
if response.status_code == 200:
return
@@ -667,22 +736,23 @@ class DingTalkClient:
'DingTalk createAndDeliver request body: %s',
json.dumps(body, ensure_ascii=False)[:1500],
)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, json=body, timeout=30.0)
response_body = await httpclient.response_text(response, max_chars=500)
if response.status_code == 200:
_stdout_logger.info(
'DingTalk createAndDeliver response: %s',
response.text[:500],
response_body,
)
return True
_stdout_logger.error(
'DingTalk createAndDeliver failed: status=%s body=%s',
response.status_code,
response.text,
response_body,
)
if self.logger:
await self.logger.error(
f'DingTalk createAndDeliver failed: status={response.status_code} body={response.text}'
f'DingTalk createAndDeliver failed: status={response.status_code} body={response_body}'
)
return False
except Exception:
@@ -725,13 +795,14 @@ class DingTalkClient:
'Content-Type': 'application/json',
}
try:
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.put(url, headers=headers, json=body, timeout=30.0)
if response.status_code == 200:
return True
if self.logger:
response_body = await httpclient.response_text(response)
await self.logger.error(
f'DingTalk card streaming failed: status={response.status_code} body={response.text}'
f'DingTalk card streaming failed: status={response.status_code} body={response_body}'
)
return False
except Exception:
@@ -768,18 +839,19 @@ class DingTalkClient:
out_track_id,
json.dumps(body, ensure_ascii=False)[:1500],
)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.put(url, headers=headers, json=body, timeout=30.0)
response_body = await httpclient.response_text(response, max_chars=300)
_stdout_logger.info(
'DingTalk update_card_data response: status=%d body=%s',
response.status_code,
response.text[:300],
response_body,
)
if response.status_code == 200:
return True
if self.logger:
await self.logger.error(
f'DingTalk update card failed: status={response.status_code} body={response.text}'
f'DingTalk update card failed: status={response.status_code} body={response_body}'
)
return False
except Exception:
@@ -808,17 +880,18 @@ class DingTalkClient:
url = 'https://oapi.dingtalk.com/gettoken'
try:
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.get(url, params={'appkey': self.key, 'appsecret': self.secret}, timeout=15.0)
data = response.json() if response.status_code == 200 else {}
data = await httpclient.parse_json_response(response) if response.status_code == 200 else {}
if data.get('errcode') == 0 and data.get('access_token'):
self.legacy_access_token = data['access_token']
expires_in = int(data.get('expires_in', 7200))
self.legacy_access_token_expiry_time = now + expires_in - 60
return self.legacy_access_token
if self.logger:
response_body = await httpclient.response_text(response, max_chars=200)
await self.logger.error(
f'DingTalk legacy gettoken failed: status={response.status_code} body={response.text[:200]}'
f'DingTalk legacy gettoken failed: status={response.status_code} body={response_body}'
)
except Exception:
_stdout_logger.exception('DingTalk legacy gettoken error')
@@ -848,8 +921,7 @@ class DingTalkClient:
url = 'https://oapi.dingtalk.com/media/upload'
try:
with open(file_path, 'rb') as f:
file_bytes = f.read()
file_bytes = await asyncio.to_thread(_read_local_media_limited, file_path)
file_name = os.path.basename(file_path)
# Best-effort content-type guess; DingTalk accepts the major image
# mime types and otherwise infers from the bytes.
@@ -857,20 +929,21 @@ class DingTalkClient:
mime = {'png': 'image/png', 'jpg': 'image/jpeg', 'jpeg': 'image/jpeg', 'gif': 'image/gif'}.get(
ext, 'application/octet-stream'
)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(
url,
params={'access_token': token, 'type': 'image'},
files={'media': (file_name, file_bytes, mime)},
timeout=30.0,
)
data = response.json() if response.status_code == 200 else {}
data = await httpclient.parse_json_response(response) if response.status_code == 200 else {}
if data.get('errcode') == 0 and data.get('media_id'):
_stdout_logger.info('DingTalk upload_image_media OK: media_id=%s', data['media_id'])
return data['media_id']
if self.logger:
response_body = await httpclient.response_text(response, max_chars=300)
await self.logger.error(
f'DingTalk upload_image_media failed: status={response.status_code} body={response.text[:300]}'
f'DingTalk upload_image_media failed: status={response.status_code} body={response_body}'
)
except Exception:
_stdout_logger.exception('DingTalk upload_image_media error')
@@ -897,15 +970,19 @@ class DingTalkClient:
continue
uri = '%s?ticket=%s' % (connection['endpoint'], urllib.parse.quote_plus(connection['ticket']))
async with websockets.connect(uri) as websocket:
async with websockets.connect(uri, max_size=_MAX_GATEWAY_MESSAGE_BYTES) as websocket:
self.client.websocket = websocket
keepalive_task = asyncio.create_task(self._keepalive(websocket))
try:
async for raw_message in websocket:
if self._stopped:
break
json_message = json.loads(raw_message)
asyncio.create_task(self.client.background_task(json_message))
json_message = await asyncio.to_thread(json.loads, raw_message)
if not self._start_inbound_task(self.client.background_task(json_message)):
if self.logger:
await self.logger.warning(
'DingTalk inbound task capacity reached; dropping message'
)
finally:
keepalive_task.cancel()
try:
@@ -948,5 +1025,15 @@ class DingTalkClient:
await self.client.websocket.close()
except Exception:
pass
inbound_tasks = list(self._inbound_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_tasks.clear()
# Clear message handlers to prevent stale callbacks
self._message_handlers = {'example': []}
if self._http_client is not None:
await self._http_client.aclose()
self._http_client = None
src/langbot/libs/official_account_api/api.py
+100 -9
View File
@@ -21,8 +21,14 @@ xml_template = """
</xml>
"""
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
class OAClient:
_STATE_TTL_SECONDS = 600
_STATE_MAX = 4096
_MAX_CONTENT_CHARS = 200000
def __init__(
self,
token: str,
@@ -41,6 +47,7 @@ class OAClient:
self.access_token = ''
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -57,8 +64,38 @@ class OAClient:
self.access_token_expiry_time = None
self.msg_id_map = {}
self.generated_content = {}
self._msg_seen_at = {}
self._generated_at = {}
self._last_state_prune = 0.0
self.logger = logger
def _prune_state(self) -> None:
now = time.monotonic()
if now - self._last_state_prune >= 60:
self._last_state_prune = now
for message_id, seen_at in tuple(self._msg_seen_at.items()):
if now - seen_at > self._STATE_TTL_SECONDS:
self._msg_seen_at.pop(message_id, None)
self.msg_id_map.pop(message_id, None)
for message_id, generated_at in tuple(self._generated_at.items()):
if now - generated_at > self._STATE_TTL_SECONDS:
self._generated_at.pop(message_id, None)
self.generated_content.pop(message_id, None)
while len(self.msg_id_map) > self._STATE_MAX:
message_id = next(iter(self.msg_id_map))
self.msg_id_map.pop(message_id, None)
self._msg_seen_at.pop(message_id, None)
while len(self.generated_content) > self._STATE_MAX:
message_id = next(iter(self.generated_content))
self.generated_content.pop(message_id, None)
self._generated_at.pop(message_id, None)
def clear(self) -> None:
self.msg_id_map.clear()
self.generated_content.clear()
self._msg_seen_at.clear()
self._generated_at.clear()
async def handle_callback_request(self):
"""处理回调请求（独立端口模式，使用全局 request）。"""
return await self._handle_callback_internal(request)
@@ -104,8 +141,16 @@ class OAClient:
raise Exception('拒绝请求')
elif req.method == 'POST':
encryt_msg = await req.data
if len(encryt_msg) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('Official Account callback body exceeds the size limit')
wxcpt = WXBizMsgCrypt(self.token, self.aes, self.appid)
ret, xml_msg = wxcpt.DecryptMsg(encryt_msg, msg_signature, timestamp, nonce)
ret, xml_msg = await asyncio.to_thread(
wxcpt.DecryptMsg,
encryt_msg,
msg_signature,
timestamp,
nonce,
)
xml_msg = xml_msg.decode('utf-8')
if ret != 0:
@@ -118,7 +163,7 @@ class OAClient:
if event:
await self._handle_message(event)
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
from_user = root.find('FromUserName').text # 发送者
to_user = root.find('ToUserName').text # 机器人
@@ -126,6 +171,7 @@ class OAClient:
interval = 0.1
while True:
content = self.generated_content.pop(message_data['MsgId'], None)
self._generated_at.pop(message_data['MsgId'], None)
if content:
response_xml = xml_template.format(
to_user=from_user,
@@ -156,7 +202,7 @@ class OAClient:
traceback.print_exc()
async def get_message(self, xml_msg: str):
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
message_data = {
'ToUserName': root.find('ToUserName').text,
@@ -193,21 +239,30 @@ class OAClient:
处理消息事件。
"""
message_id = event.message_id
self._prune_state()
if message_id in self.msg_id_map.keys():
self.msg_id_map[message_id] += 1
self._msg_seen_at[message_id] = time.monotonic()
return
self.msg_id_map[message_id] = 1
self._msg_seen_at[message_id] = time.monotonic()
msg_type = event.type
if msg_type in self._message_handlers:
for handler in self._message_handlers[msg_type]:
await handler(event)
async def set_message(self, msg_id: int, content: str):
self.generated_content[msg_id] = content
self.generated_content[msg_id] = str(content)[: self._MAX_CONTENT_CHARS]
self._generated_at[msg_id] = time.monotonic()
self._prune_state()
class OAClientForLongerResponse:
_MAX_USERS = 4096
_MAX_MESSAGES_PER_USER = 20
_MAX_CONTENT_CHARS = 200000
def __init__(
self,
token: str,
@@ -227,6 +282,7 @@ class OAClientForLongerResponse:
self.access_token = ''
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -244,8 +300,28 @@ class OAClientForLongerResponse:
self.loading_message = LoadingMessage
self.msg_queue = {}
self.user_msg_queue = {}
self._last_queue_cleanup = 0.0
self.logger = logger
def _prune_queues(self) -> None:
now = time.monotonic()
if now - self._last_queue_cleanup >= 60:
self._last_queue_cleanup = now
for user_id, queue in tuple(self.msg_queue.items()):
if not queue:
self.msg_queue.pop(user_id, None)
for user_id, queue in tuple(self.user_msg_queue.items()):
if not queue:
self.user_msg_queue.pop(user_id, None)
while len(self.msg_queue) > self._MAX_USERS:
self.msg_queue.pop(next(iter(self.msg_queue)), None)
while len(self.user_msg_queue) > self._MAX_USERS:
self.user_msg_queue.pop(next(iter(self.user_msg_queue)), None)
def clear(self) -> None:
self.msg_queue.clear()
self.user_msg_queue.clear()
async def handle_callback_request(self):
"""处理回调请求（独立端口模式，使用全局 request）。"""
return await self._handle_callback_internal(request)
@@ -285,8 +361,16 @@ class OAClientForLongerResponse:
elif req.method == 'POST':
encryt_msg = await req.data
if len(encryt_msg) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('Official Account callback body exceeds the size limit')
wxcpt = WXBizMsgCrypt(self.token, self.aes, self.appid)
ret, xml_msg = wxcpt.DecryptMsg(encryt_msg, msg_signature, timestamp, nonce)
ret, xml_msg = await asyncio.to_thread(
wxcpt.DecryptMsg,
encryt_msg,
msg_signature,
timestamp,
nonce,
)
xml_msg = xml_msg.decode('utf-8')
if ret != 0:
@@ -294,7 +378,7 @@ class OAClientForLongerResponse:
raise Exception('消息解密失败')
# 解析 XML
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
from_user = root.find('FromUserName').text
to_user = root.find('ToUserName').text
@@ -305,6 +389,7 @@ class OAClientForLongerResponse:
# 弹出用户消息
if self.user_msg_queue.get(from_user) and self.user_msg_queue[from_user]:
self.user_msg_queue[from_user].pop(0)
self._prune_queues()
response_xml = xml_template.format(
to_user=from_user,
@@ -332,9 +417,13 @@ class OAClientForLongerResponse:
if event:
self.user_msg_queue.setdefault(from_user, []).append(
{
'content': event.message,
'content': str(event.message)[: self._MAX_CONTENT_CHARS],
}
)
self.user_msg_queue[from_user] = self.user_msg_queue[from_user][
-self._MAX_MESSAGES_PER_USER :
]
self._prune_queues()
await self._handle_message(event)
return response_xml
@@ -344,7 +433,7 @@ class OAClientForLongerResponse:
traceback.print_exc()
async def get_message(self, xml_msg: str):
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
message_data = {
'ToUserName': root.find('ToUserName').text,
@@ -393,6 +482,8 @@ class OAClientForLongerResponse:
self.msg_queue[from_user].append(
{
'msg_id': message_id,
'content': content,
'content': str(content)[: self._MAX_CONTENT_CHARS],
}
)
self.msg_queue[from_user] = self.msg_queue[from_user][-self._MAX_MESSAGES_PER_USER :]
self._prune_queues()
src/langbot/libs/openclaw_weixin_api/client.py
+70 -42
View File
@@ -10,7 +10,9 @@ from __future__ import annotations
import asyncio
import base64
import hashlib
import io
import json
import logging
import os
import struct
@@ -21,6 +23,8 @@ from urllib.parse import quote
import aiohttp
from langbot.pkg.utils import httpclient
from .types import (
ApiError,
CDNMedia,
@@ -58,6 +62,51 @@ DEFAULT_BOT_TYPE = '3'
# Maximum text length per message chunk (WeChat limit)
MAX_TEXT_CHUNK_SIZE = 2000
MAX_CDN_MEDIA_BYTES = 10 * 1024 * 1024
async def _response_text(response: aiohttp.ClientResponse) -> str:
body = await httpclient.read_limited(
response,
max_bytes=MAX_CDN_MEDIA_BYTES,
)
return body.decode('utf-8', errors='replace')
async def _response_json(response: aiohttp.ClientResponse) -> dict:
payload = json.loads(await _response_text(response))
if not isinstance(payload, dict):
raise ApiError('OpenClaw API returned a non-object response', status=response.status)
return payload
def _decrypt_cdn_payload(encrypted: bytes, aes_key: bytes) -> bytes:
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
cipher = Cipher(algorithms.AES(aes_key), modes.ECB())
decryptor = cipher.decryptor()
padded = decryptor.update(encrypted) + decryptor.finalize()
unpadder = PKCS7(128).unpadder()
return unpadder.update(padded) + unpadder.finalize()
def _encrypt_cdn_payload(
file_bytes: bytes,
) -> tuple[str, str, bytes, str]:
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
raw_key = os.urandom(16)
aes_key_hex = raw_key.hex()
encoded_key = base64.b64encode(aes_key_hex.encode('utf-8')).decode('utf-8')
padder = PKCS7(128).padder()
padded = padder.update(file_bytes) + padder.finalize()
cipher = Cipher(algorithms.AES(raw_key), modes.ECB())
encryptor = cipher.encryptor()
encrypted = encryptor.update(padded) + encryptor.finalize()
raw_md5 = hashlib.md5(file_bytes).hexdigest()
return aes_key_hex, encoded_key, encrypted, raw_md5
def _random_wechat_uin() -> str:
@@ -125,12 +174,12 @@ class OpenClawWeixinClient:
url, json=payload, headers=headers, timeout=aiohttp.ClientTimeout(total=timeout)
) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(
f'OpenClaw API error {resp.status}: {text}',
status=resp.status,
)
data = await resp.json(content_type=None)
data = await _response_json(resp)
# Check for application-level errors in the response body
errcode = data.get('errcode') or data.get('ret')
@@ -170,12 +219,12 @@ class OpenClawWeixinClient:
timeout=aiohttp.ClientTimeout(total=timeout),
) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(
f'OpenClaw API error {resp.status}: {text}',
status=resp.status,
)
data = await resp.json(content_type=None)
data = await _response_json(resp)
except (asyncio.TimeoutError, aiohttp.ServerTimeoutError):
return GetUpdatesResponse(ret=0, msgs=[], get_updates_buf=get_updates_buf)
@@ -258,9 +307,6 @@ class OpenClawWeixinClient:
Returns:
Decrypted file bytes.
"""
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
if not media.encrypt_query_param:
raise ApiError('CDN media has no encrypt_query_param', status=0)
if not media.aes_key:
@@ -285,17 +331,14 @@ class OpenClawWeixinClient:
async with session.get(cdn_url, timeout=aiohttp.ClientTimeout(total=120)) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(f'CDN download failed: {resp.status} {text}', status=resp.status)
encrypted = await resp.read()
encrypted = await httpclient.read_limited(
resp,
max_bytes=MAX_CDN_MEDIA_BYTES,
)
# Decrypt AES-128-ECB with PKCS7 padding
cipher = Cipher(algorithms.AES(aes_key), modes.ECB())
decryptor = cipher.decryptor()
padded = decryptor.update(encrypted) + decryptor.finalize()
unpadder = PKCS7(128).unpadder()
return unpadder.update(padded) + unpadder.finalize()
return await asyncio.to_thread(_decrypt_cdn_payload, encrypted, aes_key)
async def upload_media(
self,
@@ -313,28 +356,13 @@ class OpenClawWeixinClient:
Returns:
CDNMedia with encrypt_query_param and aes_key for use in sendMessage.
"""
import hashlib
if len(file_bytes) > MAX_CDN_MEDIA_BYTES:
raise ApiError('CDN media exceeds the size limit', status=0)
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
# 1. Generate random 16-byte AES key
raw_key = os.urandom(16)
aes_key_hex = raw_key.hex() # 32-char hex string
# 2. Encode key for CDNMedia: base64(hex_string) — same for all media types
# Matches official SDK: Buffer.from(aeskey_hex).toString("base64")
encoded_key = base64.b64encode(aes_key_hex.encode('utf-8')).decode('utf-8')
# 3. Encrypt file with AES-128-ECB + PKCS7
padder = PKCS7(128).padder()
padded = padder.update(file_bytes) + padder.finalize()
cipher = Cipher(algorithms.AES(raw_key), modes.ECB())
encryptor = cipher.encryptor()
encrypted = encryptor.update(padded) + encryptor.finalize()
# 4. Get upload URL
raw_md5 = hashlib.md5(file_bytes).hexdigest()
aes_key_hex, encoded_key, encrypted, raw_md5 = await asyncio.to_thread(
_encrypt_cdn_payload,
file_bytes,
)
filekey = os.urandom(16).hex() # 32-char hex, matches official SDK
upload_resp = await self.get_upload_url(
@@ -370,7 +398,7 @@ class OpenClawWeixinClient:
timeout=aiohttp.ClientTimeout(total=120),
) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
logger.error('CDN upload failed: status=%d url=%s body=%s', resp.status, cdn_url, text[:500])
raise ApiError(f'CDN upload failed: {resp.status} {text}', status=resp.status)
download_param = resp.headers.get('x-encrypted-param', '')
@@ -491,12 +519,12 @@ class OpenClawWeixinClient:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=DEFAULT_API_TIMEOUT)) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(
f'Failed to fetch QR code: {resp.status} {text}',
status=resp.status,
)
data = await resp.json(content_type=None)
data = await _response_json(resp)
logger.debug(
'fetch_qrcode response: qrcode=%s, img=%s', data.get('qrcode'), bool(data.get('qrcode_img_content'))
@@ -536,12 +564,12 @@ class OpenClawWeixinClient:
url, headers=headers, timeout=aiohttp.ClientTimeout(total=DEFAULT_QR_POLL_TIMEOUT)
) as resp:
if resp.status != 200:
text = await resp.text()
text = await _response_text(resp)
raise ApiError(
f'Failed to poll QR status: {resp.status} {text}',
status=resp.status,
)
data = await resp.json(content_type=None)
data = await _response_json(resp)
logger.debug('QR status poll response: %s', data)
except (asyncio.TimeoutError, aiohttp.ServerTimeoutError):
return QRStatusResponse(status='wait')
src/langbot/libs/qq_official_api/api.py
+79 -36
View File
@@ -9,10 +9,13 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events
from .qqofficialevent import QQOfficialEvent
import json
import traceback
from contextlib import asynccontextmanager
from cryptography.hazmat.primitives.asymmetric import ed25519
from langbot.pkg.utils import httpclient
QQ_SELECT_ACTION_PREFIX = '__langbot_select__:'
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
def get_select_field_options(form_data: dict) -> tuple[str, list[str]]:
@@ -152,6 +155,7 @@ class QQOfficialClient:
def __init__(self, secret: str, token: str, app_id: str, logger: None, unified_mode: bool = False):
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -176,6 +180,32 @@ class QQOfficialClient:
self.logger = logger
self._msg_seq_counter = 0
self._token_refresh_task: Optional[asyncio.Task] = None
self._http_clients: dict[float | None, httpx.AsyncClient] = {}
@asynccontextmanager
async def _http_client_context(self, timeout: float | None = None):
client = self._http_clients.get(timeout)
if client is None or client.is_closed:
response_hooks = httpclient.httpx_response_limit_hooks()
client = (
httpx.AsyncClient(event_hooks=response_hooks)
if timeout is None
else httpx.AsyncClient(timeout=timeout, event_hooks=response_hooks)
)
self._http_clients[timeout] = client
yield client
async def close(self) -> None:
"""Stop client-owned background work."""
if self._token_refresh_task and not self._token_refresh_task.done():
self._token_refresh_task.cancel()
await asyncio.gather(self._token_refresh_task, return_exceptions=True)
self._token_refresh_task = None
clients = list(self._http_clients.values())
self._http_clients.clear()
if clients:
await asyncio.gather(*(client.aclose() for client in clients), return_exceptions=True)
async def check_access_token(self):
"""检查access_token是否存在"""
@@ -186,7 +216,7 @@ class QQOfficialClient:
async def get_access_token(self):
"""获取access_token"""
url = 'https://bots.qq.com/app/getAppAccessToken'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'appId': self.app_id,
'clientSecret': self.secret,
@@ -196,8 +226,9 @@ class QQOfficialClient:
}
response = await client.post(url, json=params, headers=headers)
if response.status_code != 200:
raise Exception(f'Failed to get access_token: HTTP {response.status_code} {response.text}')
response_data = response.json()
body = await httpclient.response_text(response)
raise Exception(f'Failed to get access_token: HTTP {response.status_code} {body}')
response_data = await httpclient.parse_json_response(response)
access_token = response_data.get('access_token')
expires_in = int(response_data.get('expires_in', 7200))
self.access_token_expiry_time = time.time() + expires_in - 60
@@ -236,8 +267,10 @@ class QQOfficialClient:
if not body or len(body) == 0:
await self.logger.info('Received empty body, might be health check or GET request')
return {'code': 0, 'message': 'ok'}, 200
if len(body) > _MAX_CALLBACK_BODY_BYTES:
return {'error': 'callback body exceeds the size limit'}, 413
payload = json.loads(body)
payload = await asyncio.to_thread(json.loads, body)
if payload.get('op') == 13:
validation_data = payload.get('d')
@@ -367,7 +400,7 @@ class QQOfficialClient:
await self.get_access_token()
url = self.base_url + '/v2/users/' + user_openid + '/messages'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -382,7 +415,7 @@ class QQOfficialClient:
if event_id:
data['event_id'] = event_id
response = await client.post(url, headers=headers, json=data)
response_data = response.json()
response_data = await httpclient.parse_json_response(response)
if response.status_code == 200:
return
else:
@@ -406,7 +439,7 @@ class QQOfficialClient:
await self.get_access_token()
url = self.base_url + '/v2/groups/' + group_openid + '/messages'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -424,8 +457,9 @@ class QQOfficialClient:
if response.status_code == 200:
return
else:
await self.logger.error(f'Failed to send group message: {response.json()}')
raise Exception(response.read().decode())
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'Failed to send group message: {error_payload}')
raise Exception(str(error_payload))
async def send_channle_group_text_msg(self, channel_id: str, content: str, msg_id: str):
"""发送频道群聊消息"""
@@ -433,7 +467,7 @@ class QQOfficialClient:
await self.get_access_token()
url = self.base_url + '/channels/' + channel_id + '/messages'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -447,7 +481,8 @@ class QQOfficialClient:
if response.status_code == 200:
return True
else:
await self.logger.error(f'Failed to send channel group message: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'Failed to send channel group message: {error_payload}')
raise Exception(response)
async def send_channle_private_text_msg(self, guild_id: str, content: str, msg_id: str):
@@ -456,7 +491,7 @@ class QQOfficialClient:
await self.get_access_token()
url = self.base_url + '/dms/' + guild_id + '/messages'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -470,7 +505,8 @@ class QQOfficialClient:
if response.status_code == 200:
return True
else:
await self.logger.error(f'Failed to send channel private message: {response.json()}')
error_payload = await httpclient.parse_json_response(response)
await self.logger.error(f'Failed to send channel private message: {error_payload}')
raise Exception(response)
# ---- 富媒体消息 ----
@@ -532,20 +568,21 @@ class QQOfficialClient:
if file_type == self.MEDIA_TYPE_FILE and file_name:
body['file_name'] = file_name
async with httpx.AsyncClient(timeout=120) as client:
async with self._http_client_context(timeout=120) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
}
response = await client.post(url, headers=headers, json=body)
if response.status_code == 200:
data = response.json()
data = await httpclient.parse_json_response(response)
file_info = data.get('file_info', '')
preview = file_info[:80] + '...' if len(file_info) > 80 else file_info
await self.logger.info(f'Upload media success, file_info={preview}')
return file_info
else:
raise Exception(f'Failed to upload media: HTTP {response.status_code} {response.text}')
body = await httpclient.response_text(response)
raise Exception(f'Failed to upload media: HTTP {response.status_code} {body}')
async def _send_media_msg(
self,
@@ -578,7 +615,7 @@ class QQOfficialClient:
if msg_id:
body['msg_id'] = msg_id
async with httpx.AsyncClient(timeout=120) as client:
async with self._http_client_context(timeout=120) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -586,7 +623,8 @@ class QQOfficialClient:
await self.logger.info(f'Sending rich media: {json.dumps(body, ensure_ascii=False)[:200]}')
response = await client.post(url, headers=headers, json=body)
if response.status_code != 200:
raise Exception(f'Failed to send rich media message: HTTP {response.status_code} {response.text}')
response_body = await httpclient.response_text(response)
raise Exception(f'Failed to send rich media message: HTTP {response.status_code} {response_body}')
async def send_image_msg(
self,
@@ -678,15 +716,16 @@ class QQOfficialClient:
if stream_msg_id:
body['stream_msg_id'] = stream_msg_id
async with httpx.AsyncClient(timeout=120) as client:
async with self._http_client_context(timeout=120) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
}
response = await client.post(url, headers=headers, json=body)
if response.status_code != 200:
raise Exception(f'Failed to send stream message: HTTP {response.status_code} {response.text}')
return response.json()
response_body = await httpclient.response_text(response)
raise Exception(f'Failed to send stream message: HTTP {response.status_code} {response_body}')
return await httpclient.parse_json_response(response)
async def send_markdown_keyboard(
self,
@@ -743,18 +782,19 @@ class QQOfficialClient:
if event_id:
body['event_id'] = event_id
async with httpx.AsyncClient(timeout=30) as client:
async with self._http_client_context(timeout=30) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
}
response = await client.post(url, headers=headers, json=body)
if response.status_code != 200:
response_body = await httpclient.response_text(response)
await self.logger.error(
f'Failed to send markdown+keyboard: HTTP {response.status_code} {response.text}'
f'Failed to send markdown+keyboard: HTTP {response.status_code} {response_body}'
)
raise Exception(f'Failed to send markdown+keyboard: HTTP {response.status_code} {response.text}')
return response.json()
raise Exception(f'Failed to send markdown+keyboard: HTTP {response.status_code} {response_body}')
return await httpclient.parse_json_response(response)
async def ack_interaction(self, interaction_id: str, code: int = 0) -> None:
"""Acknowledge a button-click INTERACTION_CREATE event.
@@ -775,7 +815,7 @@ class QQOfficialClient:
await self.get_access_token()
url = f'{self.base_url}/interactions/{interaction_id}'
async with httpx.AsyncClient(timeout=10) as client:
async with self._http_client_context(timeout=10) as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
'Content-Type': 'application/json',
@@ -783,8 +823,9 @@ class QQOfficialClient:
try:
response = await client.put(url, headers=headers, json={'code': code})
if response.status_code >= 400:
response_body = await httpclient.response_text(response)
await self.logger.warning(
f'ack_interaction non-success: HTTP {response.status_code} {response.text}'
f'ack_interaction non-success: HTTP {response.status_code} {response_body}'
)
except Exception as e:
await self.logger.warning(f'ack_interaction error (non-fatal): {e}')
@@ -796,10 +837,11 @@ class QQOfficialClient:
return time.time() > self.access_token_expiry_time
async def repeat_seed(self, bot_secret: str, target_size: int = 32) -> bytes:
seed = bot_secret
while len(seed) < target_size:
seed *= 2
return seed[:target_size].encode('utf-8')
if not bot_secret:
raise ValueError('QQ bot secret must not be empty')
target_size = max(int(target_size), 1)
repeats = (target_size + len(bot_secret) - 1) // len(bot_secret)
return (bot_secret * repeats)[:target_size].encode('utf-8')
async def verify(self, validation_payload: dict):
seed = await self.repeat_seed(self.secret)
@@ -843,19 +885,20 @@ class QQOfficialClient:
await self.get_access_token()
url = f'{self.base_url}/gateway'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
headers = {
'Authorization': f'QQBot {self.access_token}',
}
response = await client.get(url, headers=headers)
if response.status_code == 200:
data = response.json()
data = await httpclient.parse_json_response(response)
ws_url = data.get('url', '')
if not ws_url:
raise Exception('Gateway URL is empty')
return ws_url
else:
raise Exception(f'Failed to get Gateway URL: HTTP {response.status_code} {response.text}')
body = await httpclient.response_text(response)
raise Exception(f'Failed to get Gateway URL: HTTP {response.status_code} {body}')
async def _background_token_refresh(self):
"""在 token 到期前主动刷新"""
@@ -935,7 +978,7 @@ class QQOfficialClient:
try:
await self.logger.info('Connecting to WebSocket gateway...')
ws = await websockets.connect(ws_url)
ws = await websockets.connect(ws_url, max_size=_MAX_CALLBACK_BODY_BYTES)
await self.logger.info('WebSocket connected')
except Exception as e:
await self.logger.error(f'WebSocket connection failed: {e}')
@@ -948,7 +991,7 @@ class QQOfficialClient:
try:
async for raw_msg in ws:
try:
payload = json.loads(raw_msg)
payload = await asyncio.to_thread(json.loads, raw_msg)
except json.JSONDecodeError:
await self.logger.error(f'Failed to parse message: {raw_msg}')
continue
src/langbot/libs/slack_api/api.py
+7 -1
View File
@@ -1,3 +1,4 @@
import asyncio
import json
import traceback
from quart import Quart, jsonify, request
@@ -6,6 +7,8 @@ from .slackevent import SlackEvent
from typing import Callable
import langbot_plugin.api.entities.builtin.platform.events as platform_events
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
class SlackClient:
def __init__(self, bot_token: str, signing_secret: str, logger: None, unified_mode: bool = False):
@@ -13,6 +16,7 @@ class SlackClient:
self.signing_secret = signing_secret
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
self.client = AsyncWebClient(self.bot_token)
# 只有在非统一模式下才注册独立路由
@@ -50,7 +54,9 @@ class SlackClient:
"""
try:
body = await req.get_data()
data = json.loads(body)
if len(body) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('Slack callback body exceeds the size limit')
data = await asyncio.to_thread(json.loads, body)
if 'type' in data:
if data['type'] == 'url_verification':
return data['challenge']
src/langbot/libs/wechatpad_api/api/downloadpai.py
+37 -11
View File
@@ -1,7 +1,32 @@
from langbot.libs.wechatpad_api.util.http_util import post_json
import httpx
import asyncio
import base64
import httpx
from langbot.libs.wechatpad_api.util.http_util import post_json
from langbot.pkg.utils import httpclient
_MAX_WECHATPAD_MEDIA_BYTES = 16 * 1024 * 1024
async def _read_media_limited(response: httpx.Response) -> bytes:
content_length = response.headers.get('content-length')
if content_length is not None:
try:
declared_size = int(content_length)
except ValueError:
declared_size = None
if declared_size is not None and declared_size > _MAX_WECHATPAD_MEDIA_BYTES:
raise RuntimeError('WeChatPad media exceeds the runtime limit')
body = bytearray()
async for chunk in response.aiter_bytes(chunk_size=64 * 1024):
body.extend(chunk)
if len(body) > _MAX_WECHATPAD_MEDIA_BYTES:
raise RuntimeError('WeChatPad media exceeds the runtime limit')
return bytes(body)
class DownloadApi:
def __init__(self, base_url, token):
@@ -19,12 +44,13 @@ class DownloadApi:
return post_json(url, token=self.token, data=json_data)
async def download_url_to_base64(self, download_url):
async with httpx.AsyncClient() as client:
response = await client.get(download_url)
if response.status_code == 200:
file_bytes = response.content
base64_str = base64.b64encode(file_bytes).decode('utf-8') # 返回字符串格式
return base64_str
else:
raise Exception('获取文件失败')
async with httpx.AsyncClient(
timeout=30,
event_hooks=httpclient.httpx_response_limit_hooks(_MAX_WECHATPAD_MEDIA_BYTES),
) as client:
async with client.stream('GET', download_url) as response:
if response.status_code != 200:
raise RuntimeError('获取文件失败')
file_bytes = await _read_media_limited(response)
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
return encoded.decode('utf-8')
src/langbot/libs/wechatpad_api/util/http_util.py
+48 -11
View File
@@ -1,6 +1,29 @@
import json as json_module
import requests
from langbot.pkg.utils import httpclient
_MAX_WECHATPAD_RESPONSE_BYTES = 16 * 1024 * 1024
def _read_requests_response_limited(response: requests.Response) -> dict:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > _MAX_WECHATPAD_RESPONSE_BYTES:
raise RuntimeError('WeChatPad response exceeds the runtime limit')
except (TypeError, ValueError):
pass
body = bytearray()
for chunk in response.iter_content(chunk_size=64 * 1024):
body.extend(chunk)
if len(body) > _MAX_WECHATPAD_RESPONSE_BYTES:
raise RuntimeError('WeChatPad response exceeds the runtime limit')
result = json_module.loads(body)
if not isinstance(result, dict):
raise RuntimeError('WeChatPad returned a non-object response')
return result
def post_json(base_url, token, data=None):
headers = {'Content-Type': 'application/json'}
@@ -8,16 +31,21 @@ def post_json(base_url, token, data=None):
url = base_url + f'?key={token}'
try:
response = requests.post(url, json=data, headers=headers, timeout=60)
response.raise_for_status()
result = response.json()
with requests.post(
url,
json=data,
headers=headers,
timeout=60,
stream=True,
) as response:
response.raise_for_status()
result = _read_requests_response_limited(response)
if result:
return result
else:
raise RuntimeError(response.text)
raise RuntimeError('WeChatPad returned an empty response')
except Exception as e:
print(f'http请求失败, url={url}, exception={e}')
raise RuntimeError(str(e))
@@ -27,16 +55,20 @@ def get_json(base_url, token):
url = base_url + f'?key={token}'
try:
response = requests.get(url, headers=headers, timeout=60)
response.raise_for_status()
result = response.json()
with requests.get(
url,
headers=headers,
timeout=60,
stream=True,
) as response:
response.raise_for_status()
result = _read_requests_response_limited(response)
if result:
return result
else:
raise RuntimeError(response.text)
raise RuntimeError('WeChatPad returned an empty response')
except Exception as e:
print(f'http请求失败, url={url}, exception={e}')
raise RuntimeError(str(e))
@@ -68,7 +100,12 @@ async def async_request(
method=method, url=url, params=params, headers=headers, data=data, json=json
) as response:
response.raise_for_status() # 如果状态码不是200，抛出异常
result = await response.json()
result = json_module.loads(
await httpclient.read_limited(
response,
max_bytes=_MAX_WECHATPAD_RESPONSE_BYTES,
)
)
# print(result)
return result
# if result.get('Code') == 200:
src/langbot/libs/wecom_ai_bot_api/api.py
+124 -27
View File
@@ -10,13 +10,16 @@ import re
from typing import Any, Callable, Optional, Tuple
from urllib.parse import unquote
import httpx
from Crypto.Cipher import AES
from quart import Quart, request, Response, jsonify
from langbot.libs.wecom_ai_bot_api import wecombotevent
from langbot.libs.wecom_ai_bot_api.WXBizMsgCrypt3 import WXBizMsgCrypt
from langbot.pkg.platform.logger import EventLogger
from langbot.pkg.utils import httpclient
_CLIENT_TRANSIENT_CACHE_MAX = 4096
_MAX_STREAM_CONTENT_CHARS = 200000
@dataclass
@@ -56,7 +59,7 @@ class StreamSession:
last_access: float = field(default_factory=time.time)
# 将流水线增量结果缓存到队列，刷新请求逐条消费
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=1))
# 是否已经完成（收到最终片段）
finished: bool = False
@@ -85,6 +88,7 @@ class StreamSessionManager:
# full like → cancel → dislike feedback flow. Must align with the adapter's
# _stream_to_monitoring_msg TTL (wecombot.py).
_FEEDBACK_SESSION_TTL = 600 # 10 minutes
_MAX_SESSIONS = 4096
def __init__(self, logger: EventLogger, ttl: int = 60) -> None:
self.logger = logger
@@ -165,6 +169,26 @@ class StreamSessionManager:
if task_id:
self._task_index.pop(task_id, None)
def clear(self) -> None:
"""Release every retained stream and reverse index."""
self._sessions.clear()
self._msg_index.clear()
self._feedback_index.clear()
self._task_index.clear()
def _drop_session(self, stream_id: str) -> StreamSession | None:
session = self._sessions.pop(stream_id, None)
if session is None:
return None
if session.msg_id and self._msg_index.get(session.msg_id) == stream_id:
self._msg_index.pop(session.msg_id, None)
if session.feedback_id:
self._feedback_index.pop(session.feedback_id, None)
if session.pending_form_task_id:
self._task_index.pop(session.pending_form_task_id, None)
return session
def create_or_get(self, msg_json: dict[str, Any]) -> tuple[StreamSession, bool]:
"""根据企业微信回调创建或获取会话。
@@ -185,6 +209,14 @@ class StreamSessionManager:
session.last_access = time.time()
return session, False
self.cleanup()
while len(self._sessions) >= self._MAX_SESSIONS:
oldest_stream_id = min(
self._sessions,
key=lambda candidate: self._sessions[candidate].last_access,
)
self._drop_session(oldest_stream_id)
stream_id = str(uuid.uuid4())
session = StreamSession(
stream_id=stream_id,
@@ -221,8 +253,13 @@ class StreamSessionManager:
try:
session.queue.put_nowait(chunk)
except asyncio.QueueFull:
# 默认无界队列，此处兜底防御
await session.queue.put(chunk)
# Each chunk is a complete snapshot. Coalesce a slow consumer to
# the newest value instead of retaining every intermediate body.
try:
session.queue.get_nowait()
except asyncio.QueueEmpty:
pass
session.queue.put_nowait(chunk)
if chunk.is_final:
session.finished = True
@@ -265,7 +302,7 @@ class StreamSessionManager:
session.finished = True
session.last_access = time.time()
def cleanup(self) -> None:
def cleanup(self) -> list[str]:
"""定期清理过期会话，防止队列与映射无上限累积。
已注册 feedback_id 的会话使用更长的 TTL，确保用户在点赞/取消/点踩流程中
@@ -279,16 +316,14 @@ class StreamSessionManager:
if now - session.last_access > effective_ttl:
expired.append(stream_id)
removed_msg_ids: list[str] = []
for stream_id in expired:
session = self._sessions.pop(stream_id, None)
session = self._drop_session(stream_id)
if not session:
continue
msg_id = session.msg_id
if msg_id and self._msg_index.get(msg_id) == stream_id:
self._msg_index.pop(msg_id, None)
# Clean up feedback index for expired sessions
if session.feedback_id:
self._feedback_index.pop(session.feedback_id, None)
if session.msg_id:
removed_msg_ids.append(session.msg_id)
return removed_msg_ids
def _decrypt_file(encrypted_data: bytes, aes_key_str: str) -> bytes:
@@ -405,19 +440,19 @@ async def download_encrypted_file(
filename: Optional[str] = None
try:
async with httpx.AsyncClient(timeout=30.0) as client:
response = await client.get(download_url)
if response.status_code != 200:
await logger.error(f'Failed to download file (HTTP {response.status_code}): {response.text[:200]}')
client = httpclient.get_session()
async with client.get(download_url, timeout=30.0) as response:
if response.status != 200:
await logger.error(f'Failed to download file (HTTP {response.status})')
return None, None
encrypted_bytes = response.content
encrypted_bytes = await httpclient.read_limited(response)
filename = _extract_filename(response.headers.get('content-disposition', ''))
except Exception:
await logger.error(f'Failed to download file: {traceback.format_exc()}')
return None, None
try:
decrypted = _decrypt_file(encrypted_bytes, aes_key)
decrypted = await asyncio.to_thread(_decrypt_file, encrypted_bytes, aes_key)
return decrypted, filename
except Exception:
await logger.error(f'Failed to decrypt file: {traceback.format_exc()}')
@@ -466,7 +501,7 @@ async def parse_wecom_bot_message(
"""Download, decrypt, and convert to data URI for backward compatibility."""
data, _filename = await _safe_download(url, per_msg_aeskey)
if data:
return _bytes_to_data_uri(data)
return await asyncio.to_thread(_bytes_to_data_uri, data)
return None
if msg_type == 'text':
@@ -579,7 +614,10 @@ async def parse_wecom_bot_message(
if (file_data.get('filesize') or 0) <= max_inline_file_size:
file_bytes, dl_filename = await _safe_download(download_url, item_aeskey)
if file_bytes:
file_data['base64'] = _bytes_to_data_uri(file_bytes)
file_data['base64'] = await asyncio.to_thread(
_bytes_to_data_uri,
file_bytes,
)
if dl_filename and not file_data.get('filename'):
file_data['filename'] = dl_filename
files.append(file_data)
@@ -1567,6 +1605,8 @@ def build_multiple_interaction_update_card(
class WecomBotClient:
_MAX_DISPATCH_TASKS = 100
def __init__(
self,
Token: str,
@@ -1613,6 +1653,7 @@ class WecomBotClient:
self._feedback_callback: Optional[Callable] = None
self._card_action_callback: Optional[Callable] = None
self._stream_last_content: dict[str, str] = {}
self._dispatch_tasks: set[asyncio.Task] = set()
# Optional `source` block injected into every interactive template_card
# the client builds. Set via `set_card_source` from the adapter after
# reading config. Format: {icon_url, desc, desc_color}.
@@ -1695,7 +1736,12 @@ class WecomBotClient:
"""
reply_plain_str = json.dumps(payload, ensure_ascii=False)
reply_timestamp = str(int(time.time()))
ret, encrypt_text = self.wxcpt.EncryptMsg(reply_plain_str, nonce, reply_timestamp)
ret, encrypt_text = await asyncio.to_thread(
self.wxcpt.EncryptMsg,
reply_plain_str,
nonce,
reply_timestamp,
)
if ret != 0:
await self.logger.error(f'加密失败: {ret}')
return jsonify({'error': 'encrypt_failed'}), 500
@@ -1718,6 +1764,41 @@ class WecomBotClient:
except Exception:
await self.logger.error(traceback.format_exc())
def _start_dispatch_task(self, event: wecombotevent.WecomBotEvent) -> bool:
"""Start one bounded pipeline dispatch task."""
for task in tuple(self._dispatch_tasks):
if task.done():
self._dispatch_tasks.discard(task)
if len(self._dispatch_tasks) >= self._MAX_DISPATCH_TASKS:
return False
task = asyncio.create_task(self._dispatch_event(event))
self._dispatch_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
self._dispatch_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
async def close(self) -> None:
"""Cancel callbacks and release retained webhook state."""
dispatch_tasks = list(self._dispatch_tasks)
for task in dispatch_tasks:
if not task.done():
task.cancel()
if dispatch_tasks:
await asyncio.gather(*dispatch_tasks, return_exceptions=True)
self._dispatch_tasks.clear()
self.generated_content.clear()
self.msg_id_map.clear()
self._stream_last_content.clear()
self.stream_sessions.clear()
async def _handle_post_initial_response(self, msg_json: dict[str, Any], nonce: str) -> tuple[Response, int]:
"""处理企业微信首次推送的消息，返回 stream_id 并开启流水线。
@@ -1747,7 +1828,8 @@ class WecomBotClient:
await self.logger.error(traceback.format_exc())
else:
if is_new:
asyncio.create_task(self._dispatch_event(event))
if not self._start_dispatch_task(event):
await self.logger.warning('WeCom webhook dispatch capacity reached; dropping message')
payload = self._build_stream_payload(session.stream_id, '', False, feedback_id)
return await self._encrypt_and_reply(payload, nonce)
@@ -1870,7 +1952,10 @@ class WecomBotClient:
async def _handle_post_callback(self, req) -> tuple[Response, int] | Response:
"""处理企业微信的 POST 回调请求。"""
self.stream_sessions.cleanup()
for expired_msg_id in self.stream_sessions.cleanup():
self.generated_content.pop(expired_msg_id, None)
self._stream_last_content.pop(expired_msg_id, None)
self.msg_id_map.pop(expired_msg_id, None)
msg_signature = unquote(req.args.get('msg_signature', ''))
timestamp = unquote(req.args.get('timestamp', ''))
@@ -1883,12 +1968,18 @@ class WecomBotClient:
return Response('Bad Request', status=400)
xml_post_data = f'<xml><Encrypt><![CDATA[{encrypted_msg}]]></Encrypt></xml>'
ret, decrypted_xml = self.wxcpt.DecryptMsg(xml_post_data, msg_signature, timestamp, nonce)
ret, decrypted_xml = await asyncio.to_thread(
self.wxcpt.DecryptMsg,
xml_post_data,
msg_signature,
timestamp,
nonce,
)
if ret != 0:
await self.logger.error('解密失败')
return Response('解密失败', status=400)
msg_json = json.loads(decrypted_xml)
msg_json = await asyncio.to_thread(json.loads, decrypted_xml)
event_type = extract_wecom_event_type(msg_json)
@@ -2014,6 +2105,8 @@ class WecomBotClient:
self.msg_id_map[message_id] += 1
return
self.msg_id_map[message_id] = 1
while len(self.msg_id_map) > _CLIENT_TRANSIENT_CACHE_MAX:
self.msg_id_map.pop(next(iter(self.msg_id_map)), None)
msg_type = event.type
if msg_type in self._message_handlers:
for handler in self._message_handlers[msg_type]:
@@ -2047,6 +2140,8 @@ class WecomBotClient:
next_content = previous_content
else:
next_content = previous_content + content if previous_content else content
if len(next_content) > _MAX_STREAM_CONTENT_CHARS:
next_content = next_content[-_MAX_STREAM_CONTENT_CHARS:]
if not is_final and next_content == previous_content:
return True
@@ -2096,7 +2191,9 @@ class WecomBotClient:
"""
handled = await self.push_stream_chunk(msg_id, content, is_final=True)
if not handled:
self.generated_content[msg_id] = content
self.generated_content[msg_id] = content[-_MAX_STREAM_CONTENT_CHARS:]
while len(self.generated_content) > _CLIENT_TRANSIENT_CACHE_MAX:
self.generated_content.pop(next(iter(self.generated_content)), None)
def on_message(self, msg_type: str):
def decorator(func: Callable[[wecombotevent.WecomBotEvent], None]):
@@ -2119,7 +2216,7 @@ class WecomBotClient:
async def download_url_to_base64(self, download_url, encoding_aes_key):
data, _filename = await download_encrypted_file(download_url, encoding_aes_key, self.logger)
if data:
return _bytes_to_data_uri(data)
return await asyncio.to_thread(_bytes_to_data_uri, data)
return None
async def run_task(self, host: str, port: int, *args, **kwargs):
src/langbot/libs/wecom_ai_bot_api/ws_client.py
+134 -16
View File
@@ -47,6 +47,17 @@ CMD_RESPOND_WELCOME = 'aibot_respond_welcome_msg'
CMD_RESPOND_UPDATE = 'aibot_respond_update_msg'
CMD_SEND_MSG = 'aibot_send_msg'
_DEDUP_CACHE_MAX = 4096
_STREAM_CACHE_MAX = 1024
_FEEDBACK_CACHE_MAX = 4096
_PENDING_FORM_MAX = 1024
_PENDING_FORM_TTL_SECONDS = 1800
_MAX_STREAM_CONTENT_CHARS = 200000
_MAX_CALLBACK_TASKS = 100
_MAX_REPLY_WORKERS = 100
_MAX_REPLY_QUEUE_SIZE = 100
_MAX_PENDING_ACKS = 256
def _generate_req_id(prefix: str) -> str:
"""Generate a unique request ID in the format: {prefix}_{timestamp}_{random}."""
@@ -106,6 +117,7 @@ class WecomBotWsClient:
# Per-req_id serial reply queues
self._reply_queues: dict[str, asyncio.Queue] = {}
self._reply_workers: dict[str, asyncio.Task] = {}
self._callback_tasks: set[asyncio.Task] = set()
self._reply_ack_timeout = 5.0
# Stream ID tracking for WebSocket mode
@@ -135,6 +147,31 @@ class WecomBotWsClient:
# `set_card_source` from the adapter after reading config.
self.card_source: Optional[dict] = None
@staticmethod
def _cap_mapping(mapping: dict, max_entries: int) -> None:
while len(mapping) > max_entries:
mapping.pop(next(iter(mapping)), None)
def _prune_stream_state(self) -> None:
while len(self._stream_sessions) > _STREAM_CACHE_MAX:
msg_id = next(iter(self._stream_sessions))
self._stream_sessions.pop(msg_id, None)
self._stream_ids.pop(msg_id, None)
self._stream_last_content.pop(msg_id, None)
task_id = self._task_id_by_msg.pop(msg_id, None)
if task_id:
self._pending_forms_by_task.pop(task_id, None)
def _prune_pending_forms(self) -> None:
cutoff = time.monotonic() - _PENDING_FORM_TTL_SECONDS
for task_id, pending in tuple(self._pending_forms_by_task.items()):
if float(pending.get('created_at', 0.0)) <= cutoff:
self._drop_pending_form_task(task_id, pending)
while len(self._pending_forms_by_task) > _PENDING_FORM_MAX:
task_id = next(iter(self._pending_forms_by_task))
pending = self._pending_forms_by_task.get(task_id, {})
self._drop_pending_form_task(task_id, pending)
# ── Public API ──────────────────────────────────────────────────
async def connect(self):
@@ -173,17 +210,40 @@ class WecomBotWsClient:
async def disconnect(self):
"""Gracefully disconnect from the WebSocket server."""
self._running = False
heartbeat_tasks = []
if self._heartbeat_task and not self._heartbeat_task.done():
self._heartbeat_task.cancel()
for task in self._reply_workers.values():
heartbeat_tasks.append(self._heartbeat_task)
reply_workers = list(self._reply_workers.values())
for task in reply_workers:
if not task.done():
task.cancel()
callback_tasks = list(self._callback_tasks)
for task in callback_tasks:
if not task.done():
task.cancel()
shutdown_tasks = [*heartbeat_tasks, *reply_workers, *callback_tasks]
if shutdown_tasks:
await asyncio.gather(*shutdown_tasks, return_exceptions=True)
self._clear_pending_acks('Connection closed')
if self._ws and not self._ws.closed:
await self._ws.close()
self._ws = None
if self._session and not self._session.closed:
await self._session.close()
self._session = None
self._heartbeat_task = None
self._reply_queues.clear()
self._reply_workers.clear()
self._callback_tasks.clear()
self._stream_ids.clear()
self._stream_last_content.clear()
self._stream_sessions.clear()
self._feedback_sessions.clear()
self._msg_feedback_ids.clear()
self._pending_forms_by_task.clear()
self._task_id_by_msg.clear()
self._msg_id_map.clear()
def on_message(self, msg_type: str) -> Callable:
"""Decorator to register a message handler.
@@ -366,8 +426,10 @@ class WecomBotWsClient:
'chat_id': session_info.get('chat_id', ''),
'stream_id': stream_id,
'req_id': req_id,
'created_at': time.monotonic(),
}
self._task_id_by_msg[msg_id] = task_id
self._prune_pending_forms()
card_payload = build_human_input_template_card_payload(
form_data,
@@ -458,6 +520,8 @@ class WecomBotWsClient:
next_content = previous_content
else:
next_content = previous_content + content if previous_content else content
if len(next_content) > _MAX_STREAM_CONTENT_CHARS:
next_content = next_content[-_MAX_STREAM_CONTENT_CHARS:]
# Skip sending if content hasn't changed (e.g. during tool call argument streaming)
if not is_final and next_content == previous_content:
@@ -485,6 +549,8 @@ class WecomBotWsClient:
session_info = self._stream_sessions.get(msg_id)
if session_info:
self._feedback_sessions[feedback_id] = session_info
self._cap_mapping(self._feedback_sessions, _FEEDBACK_CACHE_MAX)
self._cap_mapping(self._msg_feedback_ids, _FEEDBACK_CACHE_MAX)
# WeCom replaces the displayed stream content on each refresh, so
# every frame must contain the complete snapshot, not only a delta.
@@ -516,7 +582,7 @@ class WecomBotWsClient:
self._session = aiohttp.ClientSession()
try:
self._ws = await self._session.ws_connect(self.ws_url)
self._ws = await self._session.ws_connect(self.ws_url, max_msg_size=1024 * 1024)
self._missed_pong_count = 0
self._reconnect_attempts = 0
await self.logger.info('WebSocket connected, sending auth...')
@@ -539,6 +605,8 @@ class WecomBotWsClient:
finally:
if self._heartbeat_task and not self._heartbeat_task.done():
self._heartbeat_task.cancel()
await asyncio.gather(self._heartbeat_task, return_exceptions=True)
self._heartbeat_task = None
self._clear_pending_acks('Connection closed')
finally:
if self._ws and not self._ws.closed:
@@ -565,7 +633,7 @@ class WecomBotWsClient:
try:
msg = await asyncio.wait_for(self._ws.receive(), timeout=10.0)
if msg.type in (aiohttp.WSMsgType.TEXT,):
frame = json.loads(msg.data)
frame = await asyncio.to_thread(json.loads, msg.data)
req_id = frame.get('headers', {}).get('req_id', '')
if req_id.startswith(CMD_SUBSCRIBE) and frame.get('errcode') == 0:
return True
@@ -614,7 +682,7 @@ class WecomBotWsClient:
break
if msg.type == aiohttp.WSMsgType.TEXT:
try:
frame = json.loads(msg.data)
frame = await asyncio.to_thread(json.loads, msg.data)
await self._handle_frame(frame)
except json.JSONDecodeError:
await self.logger.error(f'Failed to parse WebSocket message: {str(msg.data)[:200]}')
@@ -622,7 +690,7 @@ class WecomBotWsClient:
await self.logger.error(f'Error handling frame: {traceback.format_exc()}')
elif msg.type == aiohttp.WSMsgType.BINARY:
try:
frame = json.loads(msg.data)
frame = await asyncio.to_thread(json.loads, msg.data)
await self._handle_frame(frame)
except Exception:
await self.logger.error(f'Error handling binary frame: {traceback.format_exc()}')
@@ -638,12 +706,14 @@ class WecomBotWsClient:
# Message push
if cmd == CMD_MSG_CALLBACK:
asyncio.create_task(self._handle_message_callback(frame))
if not self._start_callback_task(self._handle_message_callback(frame)):
await self.logger.warning('WeCom WebSocket callback capacity reached; dropping message')
return
# Event push
if cmd == CMD_EVENT_CALLBACK:
asyncio.create_task(self._handle_event_callback(frame))
if not self._start_callback_task(self._handle_event_callback(frame)):
await self.logger.warning('WeCom WebSocket callback capacity reached; dropping event')
return
# No cmd → response/ACK frame, dispatch by req_id prefix
@@ -665,6 +735,27 @@ class WecomBotWsClient:
# Unknown frame
await self.logger.warning(f'Unknown frame: {_frame_snippet(frame)}')
def _start_callback_task(self, coro) -> bool:
"""Start one bounded inbound frame callback."""
for task in tuple(self._callback_tasks):
if task.done():
self._callback_tasks.discard(task)
if len(self._callback_tasks) >= _MAX_CALLBACK_TASKS:
coro.close()
return False
task = asyncio.create_task(coro)
self._callback_tasks.add(task)
def done(done_task: asyncio.Task) -> None:
self._callback_tasks.discard(done_task)
if not done_task.cancelled():
done_task.exception()
task.add_done_callback(done)
return True
async def _handle_message_callback(self, frame: dict):
"""Handle an incoming message callback frame."""
try:
@@ -697,6 +788,7 @@ class WecomBotWsClient:
'chat_id': message_data.get('chatid', ''),
'chat_type': message_data.get('type', 'single'),
}
self._prune_stream_state()
message_data['stream_id'] = stream_id
message_data['req_id'] = req_id
@@ -748,7 +840,7 @@ class WecomBotWsClient:
)
# Look up session by feedback_id
session_info = self._feedback_sessions.get(feedback_id)
session_info = self._feedback_sessions.pop(feedback_id, None)
session = None
if session_info:
session = StreamSession(
@@ -806,6 +898,10 @@ class WecomBotWsClient:
if pending is None:
await self.logger.warning(f'No pending_form found for task_id={task_id} (ws); card event ignored')
return
if time.monotonic() - float(pending.get('created_at', 0.0)) > _PENDING_FORM_TTL_SECONDS:
self._drop_pending_form_task(task_id, pending)
await self.logger.warning(f'Pending form expired for task_id={task_id} (ws)')
return
req_id_for_update = frame.get('headers', {}).get('req_id', '')
form_data = pending.get('form_data', {}) or {}
@@ -868,6 +964,7 @@ class WecomBotWsClient:
self._msg_id_map[message_id] += 1
return
self._msg_id_map[message_id] = 1
self._cap_mapping(self._msg_id_map, _DEDUP_CACHE_MAX)
msg_type = event.type
if msg_type in self._message_handlers:
@@ -899,40 +996,61 @@ class WecomBotWsClient:
# Ensure serial delivery per req_id
if req_id not in self._reply_queues:
self._reply_queues[req_id] = asyncio.Queue()
if len(self._reply_queues) >= _MAX_REPLY_WORKERS:
await self.logger.warning('WeCom WebSocket reply worker capacity reached; dropping reply')
return None
self._reply_queues[req_id] = asyncio.Queue(maxsize=_MAX_REPLY_QUEUE_SIZE)
self._reply_workers[req_id] = asyncio.create_task(self._reply_queue_worker(req_id))
future: asyncio.Future = asyncio.get_event_loop().create_future()
await self._reply_queues[req_id].put((frame, future))
try:
self._reply_queues[req_id].put_nowait((frame, future))
except asyncio.QueueFull:
await self.logger.warning(f'WeCom WebSocket reply queue full for req_id={req_id}; dropping reply')
return None
return await future
async def _reply_queue_worker(self, req_id: str):
"""Process reply queue items serially for a given req_id."""
queue = self._reply_queues[req_id]
current_future: asyncio.Future | None = None
try:
while self._running:
try:
frame, future = await asyncio.wait_for(queue.get(), timeout=60.0)
frame, current_future = await asyncio.wait_for(queue.get(), timeout=60.0)
except asyncio.TimeoutError:
# Queue idle, clean up worker
break
try:
ack = await self._send_and_wait_ack(frame)
if not future.done():
future.set_result(ack)
if not current_future.done():
current_future.set_result(ack)
except Exception as e:
if not future.done():
future.set_exception(e)
if not current_future.done():
current_future.set_exception(e)
finally:
current_future = None
except asyncio.CancelledError:
pass
if current_future is not None and not current_future.done():
current_future.set_exception(ConnectionError('Connection closed'))
finally:
while True:
try:
_, future = queue.get_nowait()
except asyncio.QueueEmpty:
break
if not future.done():
future.set_exception(ConnectionError('Reply worker stopped'))
self._reply_queues.pop(req_id, None)
self._reply_workers.pop(req_id, None)
async def _send_and_wait_ack(self, frame: dict) -> Optional[dict]:
"""Send a frame and wait for the corresponding ACK."""
req_id = frame['headers']['req_id']
if len(self._pending_acks) >= _MAX_PENDING_ACKS:
await self.logger.warning('WeCom WebSocket pending ACK capacity reached; dropping frame')
return None
ack_future: asyncio.Future = asyncio.get_event_loop().create_future()
self._pending_acks[req_id] = ack_future
src/langbot/libs/wecom_api/api.py
+150 -41
View File
@@ -1,16 +1,82 @@
from quart import request
from .WXBizMsgCrypt3 import WXBizMsgCrypt
import asyncio
import base64
import binascii
import contextvars
import functools
import httpx
import os
import traceback
from urllib.parse import quote
from quart import Quart
import xml.etree.ElementTree as ET
from contextlib import asynccontextmanager
from typing import Callable, Dict, Any
from .wecomevent import WecomEvent
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import aiofiles
from langbot.pkg.utils import httpclient
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
_EXTENDED_HTTP_TIMEOUT_SECONDS = 120
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
except (TypeError, ValueError) as exc:
if 'exceeds' in str(exc):
raise
content = bytearray()
async for chunk in response.aiter_bytes():
content.extend(chunk)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
return bytes(content)
async def _read_local_media_limited(path: str) -> bytes:
if await asyncio.to_thread(os.path.getsize, path) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
async with aiofiles.open(path, 'rb') as file:
content = await file.read(_MAX_MEDIA_BYTES + 1)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
return content
async def _decode_media_base64_limited(value: str) -> bytes:
max_encoded_chars = 4 * ((_MAX_MEDIA_BYTES + 2) // 3) + 4
if len(value) > max_encoded_chars:
raise ValueError('WeCom media exceeds the size limit')
content = await asyncio.to_thread(base64.b64decode, value)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
return content
def _bounded_token_retry(method):
"""Allow one token-refresh retry without unbounded async recursion."""
depth = contextvars.ContextVar(f'{method.__name__}_token_retry_depth', default=0)
@functools.wraps(method)
async def wrapped(*args, **kwargs):
current_depth = depth.get()
if current_depth >= 2:
raise RuntimeError(f'{method.__name__} exceeded the token refresh retry limit')
token = depth.set(current_depth + 1)
try:
return await method(*args, **kwargs)
finally:
depth.reset(token)
return wrapped
class WecomClient:
@@ -36,6 +102,7 @@ class WecomClient:
self.logger = logger
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -49,6 +116,29 @@ class WecomClient:
self._message_handlers = {
'example': [],
}
self._http_clients: dict[bool, httpx.AsyncClient] = {}
@asynccontextmanager
async def _http_client_context(self, *, unbounded_timeout: bool = False):
client = self._http_clients.get(unbounded_timeout)
if client is None or client.is_closed:
response_hooks = httpclient.httpx_response_limit_hooks()
client = (
httpx.AsyncClient(
timeout=_EXTENDED_HTTP_TIMEOUT_SECONDS,
event_hooks=response_hooks,
)
if unbounded_timeout
else httpx.AsyncClient(event_hooks=response_hooks)
)
self._http_clients[unbounded_timeout] = client
yield client
async def close(self) -> None:
clients = list(self._http_clients.values())
self._http_clients.clear()
if clients:
await asyncio.gather(*(client.aclose() for client in clients), return_exceptions=True)
# access——token操作
async def check_access_token(self):
@@ -59,15 +149,16 @@ class WecomClient:
async def get_access_token(self, secret):
url = f'{self.base_url}/gettoken?corpid={self.corpid}&corpsecret={secret}'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.get(url)
data = response.json()
data = await httpclient.parse_json_response(response)
if 'access_token' in data:
return data['access_token']
else:
await self.logger.error(f'获取accesstoken失败:{response.json()}')
await self.logger.error(f'获取accesstoken失败:{data}')
raise Exception(f'未获取access token: {data}')
@_bounded_token_retry
async def get_user_info(self, userid: str) -> dict:
"""
Get user information by user ID using the application secret.
@@ -82,9 +173,9 @@ class WecomClient:
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/user/get?access_token=' + self.access_token + '&userid=' + quote(userid)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.get(url)
data = response.json()
data = await httpclient.parse_json_response(response)
if data.get('errcode') == 40014 or data.get('errcode') == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.get_user_info(userid)
@@ -98,13 +189,13 @@ class WecomClient:
self.access_token_for_contacts = await self.get_access_token(self.secret_for_contacts)
url = self.base_url + '/user/list_id?access_token=' + self.access_token_for_contacts
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'cursor': '',
'limit': 10000,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 0:
dept_users = data['dept_user']
userid = []
@@ -121,7 +212,7 @@ class WecomClient:
url = self.base_url + '/message/send?access_token=' + self.access_token_for_contacts
user_ids = await self.get_users()
user_ids_string = '|'.join(user_ids)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_ids_string,
'msgtype': 'text',
@@ -135,16 +226,17 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] != 0:
raise Exception('Failed to send message: ' + str(data))
@_bounded_token_retry
async def send_image(self, user_id: str, agent_id: int, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/message/send?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_id,
'msgtype': 'image',
@@ -158,7 +250,7 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_image(user_id, agent_id, media_id)
@@ -166,11 +258,12 @@ class WecomClient:
await self.logger.error(f'发送图片失败:{data}')
raise Exception('Failed to send image: ' + str(data))
@_bounded_token_retry
async def send_voice(self, user_id: str, agent_id: int, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/message/send?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_id,
'msgtype': 'voice',
@@ -184,7 +277,7 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_voice(user_id, agent_id, media_id)
@@ -192,11 +285,12 @@ class WecomClient:
await self.logger.error(f'发送语音失败:{data}')
raise Exception('Failed to send voice: ' + str(data))
@_bounded_token_retry
async def send_file(self, user_id: str, agent_id: int, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/message/send?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_id,
'msgtype': 'file',
@@ -210,7 +304,7 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_file(user_id, agent_id, media_id)
@@ -218,12 +312,13 @@ class WecomClient:
await self.logger.error(f'发送文件失败:{data}')
raise Exception('Failed to send file: ' + str(data))
@_bounded_token_retry
async def send_private_msg(self, user_id: str, agent_id: int, content: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/message/send?access_token=' + self.access_token
async with httpx.AsyncClient(timeout=None) as client:
async with self._http_client_context(unbounded_timeout=True) as client:
params = {
'touser': user_id,
'msgtype': 'text',
@@ -237,7 +332,7 @@ class WecomClient:
'duplicate_check_interval': 1800,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_private_msg(user_id, agent_id, content)
@@ -283,7 +378,15 @@ class WecomClient:
elif req.method == 'POST':
encrypt_msg = await req.data
ret, xml_msg = wxcpt.DecryptMsg(encrypt_msg, msg_signature, timestamp, nonce)
if len(encrypt_msg) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('WeCom callback body exceeds the size limit')
ret, xml_msg = await asyncio.to_thread(
wxcpt.DecryptMsg,
encrypt_msg,
msg_signature,
timestamp,
nonce,
)
if ret != 0:
await self.logger.error('消息解密失败')
raise Exception(f'消息解密失败，错误码: {ret}')
@@ -332,7 +435,7 @@ class WecomClient:
"""
解析微信返回的 XML 消息并转换为字典。
"""
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
message_data = {
'ToUserName': root.find('ToUserName').text,
'FromUserName': root.find('FromUserName').text,
@@ -366,6 +469,7 @@ class WecomClient:
return ext
return 'jpg' # 默认返回jpg
@_bounded_token_retry
async def upload_image_to_work(self, image: platform_message.Image):
"""
获取 media_id
@@ -379,9 +483,8 @@ class WecomClient:
# 获取文件的二进制数据
if image.path:
async with aiofiles.open(image.path, 'rb') as f:
file_bytes = await f.read()
file_name = image.path.split('/')[-1]
file_bytes = await _read_local_media_limited(image.path)
file_name = image.path.split('/')[-1]
elif image.url:
file_bytes = await self.download_media_to_bytes(image.url)
file_name = image.url.split('/')[-1]
@@ -392,7 +495,7 @@ class WecomClient:
base64_data = base64_data.split(',', 1)[1]
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
padded_base64 = base64_data + '=' * padding
file_bytes = base64.b64decode(padded_base64)
file_bytes = await _decode_media_base64_limited(padded_base64)
except binascii.Error as e:
raise ValueError(f'Invalid base64 string: {str(e)}')
else:
@@ -400,6 +503,8 @@ class WecomClient:
raise ValueError('image对象出错')
# 设置 multipart/form-data 格式的文件
if len(file_bytes) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
boundary = '-------------------------acebdf13572468'
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
body = (
@@ -413,9 +518,9 @@ class WecomClient:
)
# 上传文件
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, content=body)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
media_id = await self.upload_image_to_work(image)
@@ -426,6 +531,7 @@ class WecomClient:
media_id = data.get('media_id')
return media_id
@_bounded_token_retry
async def upload_voice_to_work(self, voice: platform_message.Voice):
"""
上传语音文件到企业微信
@@ -437,9 +543,8 @@ class WecomClient:
file_name = 'voice.mp3'
if voice.path:
async with aiofiles.open(voice.path, 'rb') as f:
file_bytes = await f.read()
file_name = voice.path.split('/')[-1]
file_bytes = await _read_local_media_limited(voice.path)
file_name = voice.path.split('/')[-1]
elif voice.url:
file_bytes = await self.download_media_to_bytes(voice.url)
file_name = voice.url.split('/')[-1]
@@ -450,13 +555,15 @@ class WecomClient:
base64_data = base64_data.split(',', 1)[1]
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
padded_base64 = base64_data + '=' * padding
file_bytes = base64.b64decode(padded_base64)
file_bytes = await _decode_media_base64_limited(padded_base64)
except binascii.Error as e:
raise ValueError(f'Invalid base64 string: {str(e)}')
else:
await self.logger.error('Voice对象出错')
raise ValueError('voice对象出错')
if len(file_bytes) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
boundary = '-------------------------acebdf13572468'
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
body = (
@@ -470,9 +577,9 @@ class WecomClient:
)
# print(body)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, content=body)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
media_id = await self.upload_voice_to_work(voice)
@@ -482,6 +589,7 @@ class WecomClient:
media_id = data.get('media_id')
return media_id
@_bounded_token_retry
async def upload_file_to_work(self, file: platform_message.File):
"""
上传文件到企业微信
@@ -492,9 +600,8 @@ class WecomClient:
file_bytes = None
file_name = 'file.txt'
if file.path:
async with aiofiles.open(file.path, 'rb') as f:
file_bytes = await f.read()
file_name = file.path.split('/')[-1]
file_bytes = await _read_local_media_limited(file.path)
file_name = file.path.split('/')[-1]
elif file.url:
file_bytes = await self.download_media_to_bytes(file.url)
file_name = file.url.split('/')[-1]
@@ -505,12 +612,14 @@ class WecomClient:
base64_data = base64_data.split(',', 1)[1]
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
padded_base64 = base64_data + '=' * padding
file_bytes = base64.b64decode(padded_base64)
file_bytes = await _decode_media_base64_limited(padded_base64)
except binascii.Error as e:
raise ValueError(f'Invalid base64 string: {str(e)}')
else:
await self.logger.error('File对象出错')
raise ValueError('file对象出错')
if len(file_bytes) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom media exceeds the size limit')
boundary = '-------------------------acebdf13572468'
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
body = (
@@ -522,9 +631,9 @@ class WecomClient:
+ file_bytes
+ f'\r\n--{boundary}--\r\n'.encode('utf-8')
)
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, content=body)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
media_id = await self.upload_file_to_work(file)
@@ -535,10 +644,10 @@ class WecomClient:
return media_id
async def download_media_to_bytes(self, url: str) -> bytes:
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status()
return response.content
async with self._http_client_context() as client:
async with client.stream('GET', url) as response:
response.raise_for_status()
return await _read_httpx_media_limited(response)
# 进行media_id的获取
async def get_media_id(self, media: platform_message.Image | platform_message.Voice | platform_message.File):
src/langbot/libs/wecom_customer_service_api/api.py
+145 -39
View File
@@ -1,8 +1,13 @@
from quart import request
from ..wecom_api.WXBizMsgCrypt3 import WXBizMsgCrypt
import asyncio
import base64
import binascii
import contextvars
import functools
import httpx
import json
import os
import traceback
from quart import Quart
import xml.etree.ElementTree as ET
@@ -11,9 +16,72 @@ from .wecomcsevent import WecomCSEvent
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import aiofiles
import time
from contextlib import asynccontextmanager
from langbot.pkg.utils import httpclient
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
content_length = response.headers.get('Content-Length')
if content_length is not None:
try:
if int(content_length) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
except (TypeError, ValueError) as exc:
if 'exceeds' in str(exc):
raise
content = bytearray()
async for chunk in response.aiter_bytes():
content.extend(chunk)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
return bytes(content)
async def _read_local_media_limited(path: str) -> bytes:
if await asyncio.to_thread(os.path.getsize, path) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
async with aiofiles.open(path, 'rb') as file:
content = await file.read(_MAX_MEDIA_BYTES + 1)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
return content
async def _decode_media_base64_limited(value: str) -> bytes:
max_encoded_chars = 4 * ((_MAX_MEDIA_BYTES + 2) // 3) + 4
if len(value) > max_encoded_chars:
raise ValueError('WeCom customer-service media exceeds the size limit')
content = await asyncio.to_thread(base64.b64decode, value)
if len(content) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
return content
def _bounded_token_retry(method):
"""Allow one token-refresh retry without unbounded async recursion."""
depth = contextvars.ContextVar(f'{method.__name__}_token_retry_depth', default=0)
@functools.wraps(method)
async def wrapped(*args, **kwargs):
current_depth = depth.get()
if current_depth >= 2:
raise RuntimeError(f'{method.__name__} exceeded the token refresh retry limit')
token = depth.set(current_depth + 1)
try:
return await method(*args, **kwargs)
finally:
depth.reset(token)
return wrapped
class WecomCSClient:
_CUSTOMER_CACHE_MAX = 4096
def __init__(
self,
corpid: str,
@@ -34,10 +102,12 @@ class WecomCSClient:
self.logger = logger
self.unified_mode = unified_mode
self.app = Quart(__name__)
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
# Customer info cache: {external_userid: (info_dict, timestamp)}
self._customer_cache: dict[str, tuple[dict, float]] = {}
self._cache_ttl = 60 # Cache TTL in seconds (1 minute)
self._customer_cache_cleanup_at = 0.0
# 只有在非统一模式下才注册独立路由
if not self.unified_mode:
@@ -48,29 +118,40 @@ class WecomCSClient:
self._message_handlers = {
'example': [],
}
self._http_client: httpx.AsyncClient | None = None
@asynccontextmanager
async def _http_client_context(self):
if self._http_client is None or self._http_client.is_closed:
self._http_client = httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks())
yield self._http_client
async def close(self) -> None:
if self._http_client is not None:
await self._http_client.aclose()
self._http_client = None
@_bounded_token_retry
async def get_pic_url(self, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = f'{self.base_url}/media/get?access_token={self.access_token}&media_id={media_id}'
async with httpx.AsyncClient() as client:
response = await client.get(url)
if response.headers.get('Content-Type', '').startswith('application/json'):
data = response.json()
if data.get('errcode') in [40014, 42001]:
self.access_token = await self.get_access_token(self.secret)
return await self.get_pic_url(media_id)
else:
async with self._http_client_context() as client:
async with client.stream('GET', url) as response:
image_bytes = await _read_httpx_media_limited(response)
content_type = response.headers.get('Content-Type', '')
if content_type.startswith('application/json'):
data = json.loads(image_bytes)
if data.get('errcode') in [40014, 42001]:
self.access_token = await self.get_access_token(self.secret)
return await self.get_pic_url(media_id)
raise Exception('Failed to get image: ' + str(data))
# 否则是图片，转成 base64
image_bytes = response.content
content_type = response.headers.get('Content-Type', '')
base64_str = base64.b64encode(image_bytes).decode('utf-8')
base64_str = f'data:{content_type};base64,{base64_str}'
return base64_str
# 否则是图片，转成 base64
base64_str = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode('utf-8')
return f'data:{content_type};base64,{base64_str}'
# access——token操作
async def check_access_token(self):
@@ -81,19 +162,20 @@ class WecomCSClient:
async def get_access_token(self, secret):
url = f'{self.base_url}/gettoken?corpid={self.corpid}&corpsecret={secret}'
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.get(url)
data = response.json()
data = await httpclient.parse_json_response(response)
if 'access_token' in data:
return data['access_token']
else:
raise Exception(f'未获取access token: {data}')
@_bounded_token_retry
async def get_detailed_message_list(self, xml_msg: str):
# 在本方法中解析消息，并且获得消息的具体内容
if isinstance(xml_msg, bytes):
xml_msg = xml_msg.decode('utf-8')
root = ET.fromstring(xml_msg)
root = await asyncio.to_thread(ET.fromstring, xml_msg)
token = root.find('Token').text
open_kfid = root.find('OpenKfId').text
@@ -106,14 +188,14 @@ class WecomCSClient:
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/kf/sync_msg?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'token': token,
'voice_format': 0,
'open_kfid': open_kfid,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.get_detailed_message_list(xml_msg)
@@ -130,11 +212,12 @@ class WecomCSClient:
# await self.change_service_status(userid=external_userid,openkfid=open_kfid,servicer=servicer)
return last_msg_data
@_bounded_token_retry
async def change_service_status(self, userid: str, openkfid: str, servicer: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/kf/service_state/get?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'open_kfid': openkfid,
'external_userid': userid,
@@ -142,18 +225,19 @@ class WecomCSClient:
'servicer_userid': servicer,
}
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.change_service_status(userid, openkfid)
return await self.change_service_status(userid, openkfid, servicer)
if data['errcode'] != 0:
raise Exception('Failed to change service status: ' + str(data))
@_bounded_token_retry
async def send_image(self, user_id: str, agent_id: int, media_id: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
url = self.base_url + '/media/upload?access_token=' + self.access_token
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
params = {
'touser': user_id,
'toparty': '',
@@ -170,7 +254,7 @@ class WecomCSClient:
}
try:
response = await client.post(url, json=params)
data = response.json()
data = await httpclient.parse_json_response(response)
except Exception as e:
raise Exception('Failed to send image: ' + str(e))
@@ -182,6 +266,7 @@ class WecomCSClient:
if data['errcode'] != 0:
raise Exception('Failed to send image: ' + str(data))
@_bounded_token_retry
async def send_text_msg(self, open_kfid: str, external_userid: str, msgid: str, content: str):
if not await self.check_access_token():
self.access_token = await self.get_access_token(self.secret)
@@ -198,10 +283,10 @@ class WecomCSClient:
},
}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, json=payload)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
return await self.send_text_msg(open_kfid, external_userid, msgid, content)
@@ -250,7 +335,15 @@ class WecomCSClient:
elif req.method == 'POST':
encrypt_msg = await req.data
ret, xml_msg = wxcpt.DecryptMsg(encrypt_msg, msg_signature, timestamp, nonce)
if len(encrypt_msg) > _MAX_CALLBACK_BODY_BYTES:
raise ValueError('WeCom customer-service callback body exceeds the size limit')
ret, xml_msg = await asyncio.to_thread(
wxcpt.DecryptMsg,
encrypt_msg,
msg_signature,
timestamp,
nonce,
)
if ret != 0:
raise Exception(f'消息解密失败，错误码: {ret}')
@@ -315,6 +408,7 @@ class WecomCSClient:
return ext
return 'jpg' # 默认返回jpg
@_bounded_token_retry
async def upload_to_work(self, image: platform_message.Image):
"""
获取 media_id
@@ -328,9 +422,8 @@ class WecomCSClient:
# 获取文件的二进制数据
if image.path:
async with aiofiles.open(image.path, 'rb') as f:
file_bytes = await f.read()
file_name = image.path.split('/')[-1]
file_bytes = await _read_local_media_limited(image.path)
file_name = image.path.split('/')[-1]
elif image.url:
file_bytes = await self.download_image_to_bytes(image.url)
file_name = image.url.split('/')[-1]
@@ -341,13 +434,15 @@ class WecomCSClient:
base64_data = base64_data.split(',', 1)[1]
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
padded_base64 = base64_data + '=' * padding
file_bytes = base64.b64decode(padded_base64)
file_bytes = await _decode_media_base64_limited(padded_base64)
except binascii.Error as e:
raise ValueError(f'Invalid base64 string: {str(e)}')
else:
raise ValueError('image对象出错')
# 设置 multipart/form-data 格式的文件
if len(file_bytes) > _MAX_MEDIA_BYTES:
raise ValueError('WeCom customer-service media exceeds the size limit')
boundary = '-------------------------acebdf13572468'
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
body = (
@@ -361,9 +456,9 @@ class WecomCSClient:
)
# 上传文件
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, headers=headers, content=body)
data = response.json()
data = await httpclient.parse_json_response(response)
if data['errcode'] == 40014 or data['errcode'] == 42001:
self.access_token = await self.get_access_token(self.secret)
media_id = await self.upload_to_work(image)
@@ -374,16 +469,17 @@ class WecomCSClient:
return media_id
async def download_image_to_bytes(self, url: str) -> bytes:
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status()
return response.content
async with self._http_client_context() as client:
async with client.stream('GET', url) as response:
response.raise_for_status()
return await _read_httpx_media_limited(response)
# 进行media_id的获取
async def get_media_id(self, image: platform_message.Image):
media_id = await self.upload_to_work(image=image)
return media_id
@_bounded_token_retry
async def get_customer_info(self, external_userid: str) -> dict | None:
"""
Get customer information by external_userid with caching.
@@ -398,6 +494,11 @@ class WecomCSClient:
"""
# Check cache first
current_time = time.time()
if current_time - self._customer_cache_cleanup_at >= 30:
self._customer_cache_cleanup_at = current_time
for user_id, (_, cached_time) in tuple(self._customer_cache.items()):
if current_time - cached_time >= self._cache_ttl:
self._customer_cache.pop(user_id, None)
if external_userid in self._customer_cache:
cached_info, cached_time = self._customer_cache[external_userid]
if current_time - cached_time < self._cache_ttl:
@@ -413,9 +514,9 @@ class WecomCSClient:
'external_userid_list': [external_userid],
}
async with httpx.AsyncClient() as client:
async with self._http_client_context() as client:
response = await client.post(url, json=payload)
data = response.json()
data = await httpclient.parse_json_response(response)
if data.get('errcode') in [40014, 42001]:
self.access_token = await self.get_access_token(self.secret)
@@ -431,5 +532,10 @@ class WecomCSClient:
customer_info = customer_list[0]
# Store in cache
self._customer_cache[external_userid] = (customer_info, current_time)
while len(self._customer_cache) > self._CUSTOMER_CACHE_MAX:
self._customer_cache.pop(next(iter(self._customer_cache)), None)
return customer_info
return None
def clear(self) -> None:
self._customer_cache.clear()
src/langbot/libs/weknora_api/client.py
+71 -35
View File
@@ -6,6 +6,56 @@ import json
from .errors import WeKnoraAPIError
_MAX_WENKORA_RESPONSE_BYTES = 1024 * 1024
_MAX_WENKORA_STREAM_BYTES = 16 * 1024 * 1024
_MAX_WENKORA_SSE_LINE_BYTES = 1024 * 1024
async def _read_limited_response(response: httpx.Response) -> bytes:
body = bytearray()
async for chunk in response.aiter_bytes(chunk_size=8192):
body.extend(chunk)
if len(body) > _MAX_WENKORA_RESPONSE_BYTES:
raise WeKnoraAPIError('WeKnora response exceeds the runtime limit')
return bytes(body)
async def _iter_sse_json(
response: httpx.Response,
) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
buffer = bytearray()
total = 0
async for chunk in response.aiter_bytes(chunk_size=8192):
total += len(chunk)
if total > _MAX_WENKORA_STREAM_BYTES:
raise WeKnoraAPIError('WeKnora stream exceeds the runtime limit')
buffer.extend(chunk)
while b'\n' in buffer:
raw_line, _, remainder = buffer.partition(b'\n')
buffer = bytearray(remainder)
if len(raw_line) > _MAX_WENKORA_SSE_LINE_BYTES:
raise WeKnoraAPIError('WeKnora SSE event exceeds the runtime limit')
line = raw_line.rstrip(b'\r').strip()
if not line.startswith(b'data:'):
continue
try:
data = json.loads(line[5:].strip())
except json.JSONDecodeError:
continue
if isinstance(data, dict):
yield data
if len(buffer) > _MAX_WENKORA_SSE_LINE_BYTES:
raise WeKnoraAPIError('WeKnora SSE event exceeds the runtime limit')
line = bytes(buffer).rstrip(b'\r').strip()
if line.startswith(b'data:'):
try:
data = json.loads(line[5:].strip())
except json.JSONDecodeError:
return
if isinstance(data, dict):
yield data
class AsyncWeKnoraClient:
"""WeKnora API 客户端"""
@@ -39,19 +89,19 @@ class AsyncWeKnoraClient:
if description:
payload['description'] = description
response = await client.post(
async with client.stream(
'POST',
'/sessions',
headers={
'X-API-Key': self.api_key,
'Content-Type': 'application/json',
},
json=payload,
)
if response.status_code not in (200, 201):
raise WeKnoraAPIError(f'{response.status_code} {response.text}')
data = response.json()
) as response:
body = await _read_limited_response(response)
if response.status_code not in (200, 201):
raise WeKnoraAPIError(f'{response.status_code} {body.decode("utf-8", errors="replace")}')
data = json.loads(body)
return data['data']['id']
async def agent_chat(
@@ -107,20 +157,13 @@ class AsyncWeKnoraClient:
},
json=payload,
) as r:
async for chunk in r.aiter_lines():
if r.status_code != 200:
raise WeKnoraAPIError(f'{r.status_code} {chunk}')
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
try:
data = json.loads(chunk[5:].strip())
except json.JSONDecodeError:
continue
yield data
# 收到 error 事件后主动结束流，避免上层未 raise 时持续等待
if data.get('response_type') == 'error':
return
if r.status_code != 200:
body = await _read_limited_response(r)
raise WeKnoraAPIError(f'{r.status_code} {body.decode("utf-8", errors="replace")}')
async for data in _iter_sse_json(r):
yield data
if data.get('response_type') == 'error':
return
async def knowledge_chat(
self,
@@ -164,17 +207,10 @@ class AsyncWeKnoraClient:
},
json=payload,
) as r:
async for chunk in r.aiter_lines():
if r.status_code != 200:
raise WeKnoraAPIError(f'{r.status_code} {chunk}')
if chunk.strip() == '':
continue
if chunk.startswith('data:'):
try:
data = json.loads(chunk[5:].strip())
except json.JSONDecodeError:
continue
yield data
# 收到 error 事件后主动结束流，避免上层未 raise 时持续等待
if data.get('response_type') == 'error':
return
if r.status_code != 200:
body = await _read_limited_response(r)
raise WeKnoraAPIError(f'{r.status_code} {body.decode("utf-8", errors="replace")}')
async for data in _iter_sse_json(r):
yield data
if data.get('response_type') == 'error':
return